# -*- coding: utf-8 -*-
"""对关键帧字幕区逐帧 OCR，输出规范要求的时间线表字幕稿。

与 v1 的差异：
- 不再做聚类取代表帧，而是对每一帧都 OCR，避免代表帧模糊/漏句导致整句丢失。
- 字幕区 OCR 用底部 18%，防止双行字幕被截断（检测仍用 15%）。
- 按时间滚动去重（相似度≥0.80 视为同一句，保留最长版本）。
- 输出格式严格为：字幕稿_日期.md，正文用时间线表（时间 | 台词 | 帧文件）。

用法：
  python subtitle_ocr_v2.py --frames <帧目录> --video <源视频说明> --out <字幕稿.md>
"""
import argparse, difflib, glob, os, re, sys
import numpy as np
from PIL import Image


def fmt_ts(t):
    m = int(t // 60)
    s = t - m * 60
    return "%02d:%06.3f" % (m, s)


def ts_from_name(name):
    m = re.search(r"_(\d+\.\d+)s\.jpg", name)
    return float(m.group(1)) if m else 0.0


def crop_subtitle(img, crop_ratio):
    """从底部裁剪指定比例的区域用于 OCR。"""
    w, h = img.size
    sh = max(1, int(round(h * crop_ratio)))
    sy = h - sh
    return img.crop((0, sy, w, h))


def has_meaningful_cjk(text):
    cjk = len(re.findall(r"[一-鿿]", text))
    # 也允许纯英文/数字的短片段，但要求至少有一些可见字符
    return cjk >= 2 or len(text.strip()) >= 8


def sim(a, b):
    return difflib.SequenceMatcher(None, a, b).ratio()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--frames", required=True, help="关键帧目录")
    ap.add_argument("--video", required=True, help="源视频说明")
    ap.add_argument("--out", required=True, help="输出字幕稿路径")
    ap.add_argument("--crop", type=float, default=0.18, help="OCR 字幕区裁剪比例")
    ap.add_argument("--dup", type=float, default=0.80, help="滚动去重相似度阈值")
    args = ap.parse_args()

    from rapidocr_onnxruntime import RapidOCR
    engine = RapidOCR()

    files = sorted(glob.glob(os.path.join(args.frames, "*.jpg")))
    print("总帧数：%d" % len(files))

    raw = []
    for f in files:
        sub = crop_subtitle(Image.open(f), args.crop)
        # 2x 放大提升 OCR 识别率
        sub = sub.resize((sub.width * 2, sub.height * 2), Image.LANCZOS)
        res = engine(np.array(sub))
        if not res or not res[0]:
            continue
        txt = " ".join([line[1] for line in res[0] if line[1].strip()]).strip()
        txt = re.sub(r"\s+", " ", txt)
        if not txt or not has_meaningful_cjk(txt):
            continue
        ts = ts_from_name(os.path.basename(f))
        raw.append((ts, txt, os.path.basename(f)))

    print("OCR 有效行：%d" % len(raw))

    # 按时间滚动去重：同一句保留最长、最完整的版本，时间取最早
    merged = []
    for ts, txt, fname in raw:
        if not merged:
            merged.append([ts, txt, fname])
            continue
        last_ts, last_txt, last_fname = merged[-1]
        if sim(txt, last_txt) >= args.dup:
            # 视为同一句，保留内容更长的版本（如果当前更长则替换文本，但保留最早时间）
            if len(txt) > len(last_txt):
                merged[-1][1] = txt
                merged[-1][2] = fname
            # 时间保持 earliest
        else:
            merged.append([ts, txt, fname])

    print("去重后：%d" % len(merged))

    # 写规范时间线表
    date_stamp = os.path.basename(args.out).replace("字幕稿_", "").replace(".md", "")
    with open(args.out, "w", encoding="utf-8", newline="\n") as fo:
        fo.write("# 字幕稿_%s\n\n" % date_stamp)
        fo.write("源视频：%s\n\n" % args.video)
        fo.write("> 说明：视频无独立字幕轨，本字幕由关键帧画面字幕区逐帧 OCR 提取并按时间排序；"
                   "相似度≥%.0f%% 的连续重复句已合并，保留最完整的一句。\n\n" % (args.dup * 100))
        fo.write("| 时间 | 台词 | 帧文件 |\n")
        fo.write("|------|------|--------|\n")
        for ts, txt, fname in merged:
            fo.write("| %s | %s | %s |\n" % (fmt_ts(ts), txt, fname))
        fo.write("\n共 %d 条字幕。\n" % len(merged))

    print("已写出 → %s" % args.out)


if __name__ == "__main__":
    main()