# -*- coding: utf-8 -*-
"""
从已有 OCR 数据重新生成 SRT，不重新调用 OCR API。
用法: python rerun_srt.py
"""
import sys, os, json
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from difflib import SequenceMatcher
from tool.log.log_manager import UnifiedLogger
from tool.getsrt import (
    analyze_subtitle_blocks, analyze_h,
    extract_highest_subtitle, check_subtitle_quality_enhanced,
    evaluate_text_quality_rules, generate_srt_from_results
)

COMPOSITE_DIR = sys.argv[1] if len(sys.argv) > 1 else "./composite/BV1aL4y1p7Bi_p11"
JSONL_PATH    = os.path.join(COMPOSITE_DIR, "ocr_result.jsonl")
META_PATH     = os.path.join(COMPOSITE_DIR, "metadata.json")
OUTPUT_SRT    = os.path.join(COMPOSITE_DIR, "BV1aL4y1p7Bi_p11.new.srt")
INTERVAL      = 0.33

logger = UnifiedLogger()
logger.reset()
logger.init_logging("rerun_log.json")
logger.start_program("rerun")

# 读取已有 OCR 数据，将旧路径修正为当前实际路径
all_pic_ocr_pair = []
with open(JSONL_PATH, 'r', encoding='utf-8') as f:
    lines = f.readlines()

i = 0
while i < len(lines):
    if i + 1 < len(lines) and lines[i].startswith("pic_path: "):
        raw_path = lines[i].strip()[len("pic_path: "):]
        filename = os.path.basename(raw_path)
        actual_path = os.path.join(COMPOSITE_DIR, filename)
        try:
            ocr_json = json.loads(lines[i + 1])
            all_pic_ocr_pair.append([actual_path, ocr_json])
            i += 2
            continue
        except json.JSONDecodeError:
            pass
    i += 1

print(f"读取 {len(all_pic_ocr_pair)} 帧的OCR数据")

# 分析字幕网格位置
best_grids = analyze_subtitle_blocks(all_pic_ocr_pair, META_PATH, logger)
h = analyze_h(all_pic_ocr_pair, META_PATH)

if not best_grids.get("composite"):
    print("ERROR: 未找到合适的字幕区域")
    sys.exit(1)

# 逐帧提取字幕
results = []
DEBUG_KEYWORD = "曹小月从角落里拿出雨伞"
for idx in range(len(all_pic_ocr_pair)):
    print(f"\r处理 {idx + 1}/{len(all_pic_ocr_pair)}", end="", flush=True)
    text, conf, flag, sim = extract_highest_subtitle(
        idx, all_pic_ocr_pair, best_grids, META_PATH, h, logger
    )
    if text and not check_subtitle_quality_enhanced(text) and conf <= 90:
        text = ""
    # 打印包含关键词的帧的原始OCR和提取结果
    if DEBUG_KEYWORD in (text or ""):
        pic_path, ocr_json = all_pic_ocr_pair[idx]
        print(f"\n[DEBUG 帧{idx}] 提取: {repr(text)}")
        for td in ocr_json['TextDetections']:
            print(f"  box: {repr(td['DetectedText'])}  conf={td['Confidence']}  watermark={td.get('is_watermark','?')}")
    results.append((idx, text or ""))

non_empty = sum(1 for _, t in results if t)
print(f"\n共 {non_empty}/{len(results)} 帧有字幕")

# 生成 SRT
generate_srt_from_results(results, OUTPUT_SRT, interval_seconds=INTERVAL)
print(f"SRT 已输出: {OUTPUT_SRT}")

# 展示问题字幕附近的结果
with open(OUTPUT_SRT, 'r', encoding='utf-8') as f:
    srt_content = f.readlines()

print("\n=== 修复验证（'曹小月从角落里拿出雨伞' 附近） ===")
found = False
for j, line in enumerate(srt_content):
    if "曹小月从角落里拿出雨伞" in line:
        start = max(0, j - 3)
        end   = min(len(srt_content), j + 10)
        print("".join(srt_content[start:end]))
        print("---")
        found = True
        break

if not found:
    print("未在输出 SRT 中找到目标字幕，请检查提取结果。")