# -*- coding: utf-8 -*-
"""
阶段 B:文本去重 + 出稿(只读缓存,可反复重跑调阈值)
=======================================================
策略:基于 difflib 文本相似度折叠连续重复段,绝不使用像素级过滤。
"""
import csv
import json
import re
from difflib import SequenceMatcher
from pathlib import Path

# ═══════════════════════════════════════════════════════════════
# 全局参数(与阶段 A 共享,放文件顶部做常量)
# ═══════════════════════════════════════════════════════════════
FPS = 1
CROP_BOTTOM_RATIO = 0.2
SIM_THRESHOLD = 0.85

# 工作目录
WORK_DIR = Path(__file__).resolve().parent
OCR_RAW_PATH = WORK_DIR / "ocr_raw.jsonl"
B_MANUSCRIPT_PATH = WORK_DIR / "B稿_v2.txt"
DEBUG_CSV_PATH = WORK_DIR / "dedup_debug.csv"
BLANK_FILTERED_PATH = WORK_DIR / "blank_filtered.txt"


def similarity(a, b):
    """计算两个文本的 difflib 相似度(0.0 ~ 1.0)"""
    if not a or not b:
        return 0.0
    return SequenceMatcher(None, a, b).ratio()


def is_blank_ocr(text):
    """基于文本的空场判定:空串 / HTML幻觉 / 不含汉字"""
    t = text.strip()
    if not t:                                  # 空串
        return True
    if '<' in t and '>' in t:                  # HTML/markup 幻觉(如 <table>)
        return True
    if not re.search(r'[\u4e00-\u9fff]', t):   # 不含任何汉字 → 非本节目字幕
        return True
    return False


# ═══════════════════════════════════════════════════════════════
# B1: 读缓存
# ═══════════════════════════════════════════════════════════════
def load_ocr_raw():
    """读 ocr_raw.jsonl,按 idx(时间)升序返回记录列表"""
    if not OCR_RAW_PATH.exists():
        raise FileNotFoundError(
            f"ocr_raw.jsonl 不存在: {OCR_RAW_PATH}\n"
            "请先跑阶段 A(stage_a_extract_ocr.py)。"
        )

    records = []
    with open(OCR_RAW_PATH, "r", encoding="utf-8") as fh:
        for line in fh:
            line = line.strip()
            if not line:
                continue
            rec = json.loads(line)
            records.append(rec)

    # 按 idx 升序
    records.sort(key=lambda r: r["idx"])
    return records


# ═══════════════════════════════════════════════════════════════
# B2: 折叠连续重复段(裁判在此,基于文本不基于像素)
# ═══════════════════════════════════════════════════════════════
def collapse_segments(records):
    """
    维护"当前字幕"(文本 + 起始 t_sec),遍历每帧:
      - 文本为空 → 空场:结束当前字幕(若有),并清空,不开新段。
      - 非空且与当前字幕文本的 difflib 相似度 ≥ SIM_THRESHOLD → 同一条,并入当前段(记下这条文本备投票)。
      - 非空且相似度 < SIM_THRESHOLD → 结束并输出当前段,用这帧开新段。
    只比时间相邻帧,只折叠连续段;绝不全局去重。

    每段最终文本 = 段内出现次数最多的 OCR 文本(多数投票,抵消个别帧抖动)。
    每段时间戳 = 段内最早 t_sec。

    返回: (segments, debug_rows)
      segments: [{"start_t": float, "text": str}, ...]
      debug_rows: [{idx, t_sec, ocr_text, decision, merged_into_start_t}, ...]
    """
    segments = []
    debug_rows = []

    if not records:
        return segments, debug_rows

    current_text = None          # 当前段内用于比对的参考文本
    current_texts = []           # 当前段内所有 OCR 文本(用于投票)
    current_start_t = None       # 当前段起始 t_sec
    current_segment_indices = [] # 当前段包含的帧号列表

    def _finalize_segment():
        """输出当前段(若存在)"""
        nonlocal current_text, current_texts, current_start_t, current_segment_indices
        if current_text is not None and current_texts:
            # 多数投票:选出现次数最多的文本
            text_counts = {}
            for t in current_texts:
                text_counts[t] = text_counts.get(t, 0) + 1
            best_text = max(text_counts, key=lambda k: text_counts[k])

            segments.append({
                "start_t": current_start_t,
                "text": best_text,
            })
        current_text = None
        current_texts = []
        current_start_t = None
        current_segment_indices = []

    for rec in records:
        ocr_text = rec["text"]

        if is_blank_ocr(ocr_text):
            # 空场:结束当前字幕(若有),并清空,不开新段
            if current_text is not None:
                _finalize_segment()
            debug_rows.append({
                "idx": rec["idx"],
                "t_sec": rec["t_sec"],
                "ocr_text": ocr_text,
                "decision": "blank",
                "merged_into_start_t": "",
            })
            continue

        if current_text is None:
            # 开新段
            current_text = ocr_text
            current_texts = [ocr_text]
            current_start_t = rec["t_sec"]
            current_segment_indices = [rec["idx"]]
            debug_rows.append({
                "idx": rec["idx"],
                "t_sec": rec["t_sec"],
                "ocr_text": ocr_text,
                "decision": "new",
                "merged_into_start_t": rec["t_sec"],
            })
        else:
            sim = similarity(current_text, ocr_text)
            if sim >= SIM_THRESHOLD:
                # 同一条,并入当前段
                current_texts.append(ocr_text)
                current_segment_indices.append(rec["idx"])
                debug_rows.append({
                    "idx": rec["idx"],
                    "t_sec": rec["t_sec"],
                    "ocr_text": ocr_text,
                    "decision": "merged",
                    "merged_into_start_t": current_start_t,
                })
            else:
                # 结束当前段,用这帧开新段
                _finalize_segment()
                current_text = ocr_text
                current_texts = [ocr_text]
                current_start_t = rec["t_sec"]
                current_segment_indices = [rec["idx"]]
                debug_rows.append({
                    "idx": rec["idx"],
                    "t_sec": rec["t_sec"],
                    "ocr_text": ocr_text,
                    "decision": "new",
                    "merged_into_start_t": rec["t_sec"],
                })

    # 遍历结束,落最后一个段
    _finalize_segment()

    return segments, debug_rows


# ═══════════════════════════════════════════════════════════════
# B3: 输出 B稿_v2.txt
# ═══════════════════════════════════════════════════════════════
def write_b_manuscript(segments):
    """
    输出 B稿_v2.txt,每行格式: [XmYs] 文本
    例: [1m18s] 我是主持人蓝皓
    """
    lines = []
    for seg in segments:
        t = seg["start_t"]
        m = int(t) // 60
        s = int(t) % 60
        cleaned = re.sub(r'^[#*\->`\s]+', '', seg['text']).strip()
        lines.append(f"[{m}m{s}s] {cleaned}")

    with open(B_MANUSCRIPT_PATH, "w", encoding="utf-8") as fh:
        fh.write("\n".join(lines) + "\n")

    print(f"[B3] B稿_v2.txt 写入: {B_MANUSCRIPT_PATH} ({len(lines)} 行)")
    return lines


# ═══════════════════════════════════════════════════════════════
# B4: 输出 dedup_debug.csv
# ═══════════════════════════════════════════════════════════════
def write_debug_csv(debug_rows):
    """
    输出 dedup_debug.csv,逐帧记:
    idx, t_sec, ocr_text, decision(new/merged/blank), merged_into_start_t
    """
    with open(DEBUG_CSV_PATH, "w", encoding="utf-8", newline="") as fh:
        writer = csv.writer(fh)
        writer.writerow([
            "idx", "t_sec", "ocr_text",
            "decision", "merged_into_start_t",
        ])
        for row in debug_rows:
            writer.writerow([
                row["idx"],
                row["t_sec"],
                row["ocr_text"],
                row["decision"],
                row["merged_into_start_t"],
            ])

    print(f"[B4] dedup_debug.csv 写入: {DEBUG_CSV_PATH} ({len(debug_rows)} 行)")


# ═══════════════════════════════════════════════════════════════
# 主流程
# ═══════════════════════════════════════════════════════════════
def main():
    print(f"[B] SIM_THRESHOLD = {SIM_THRESHOLD}")

    # B1
    records = load_ocr_raw()
    print(f"[B1] 读取 ocr_raw.jsonl,共 {len(records)} 条记录")

    # B2
    segments, debug_rows = collapse_segments(records)
    print(f"[B2] 折叠完成: {len(segments)} 个不连续字幕段")
    print(f"[B2] 调试行数: {len(debug_rows)}")

    # B3
    lines = write_b_manuscript(segments)

    # B4
    write_debug_csv(debug_rows)

    # ═══════════════════════════════════════════════════════════
    # 跑完报告
    # ═══════════════════════════════════════════════════════════
    total_frames = len(records)
    ocr_success = sum(1 for r in records if r["text"].strip())
    blank_count = sum(1 for r in records if is_blank_ocr(r["text"]))

    # 收集所有"被判空场但原文本非空"的去重文本
    blank_but_nonempty = set()
    for r in records:
        t = r["text"]
        if is_blank_ocr(t) and t.strip():
            blank_but_nonempty.add(t)

    print("\n" + "=" * 60)
    print("跑完报告")
    print("=" * 60)
    print(f"抽帧总数:           {total_frames}")
    print(f"OCR 成功数(非空):   {ocr_success}")
    print(f"空场数:             {blank_count}")
    print(f"B稿_v2.txt 最终行数: {len(lines)}")
    print("=" * 60)
    print(f"B 稿路径: {B_MANUSCRIPT_PATH}")
    print(f"调试表路径: {DEBUG_CSV_PATH}")

    # 写出"被判空场但原文本非空"去重清单到文件
    blank_list = sorted(blank_but_nonempty)
    with open(BLANK_FILTERED_PATH, "w", encoding="utf-8") as fh:
        for t in blank_list:
            fh.write(t + "\n")
    print(f"\n空白过滤清单写入: {BLANK_FILTERED_PATH} ({len(blank_list)} 条去重文本)")

    # 打印"被判空场但原文本非空"的去重清单
    if blank_but_nonempty:
        print("\n" + "=" * 60)
        print("【被判空场但原文本非空】去重文本清单(请肉眼确认无真中文字幕)")
        print("=" * 60)
        for i, t in enumerate(blank_list, 1):
            print(f"  [{i}] {repr(t)}")
        print(f"共 {len(blank_but_nonempty)} 条去重文本")
    else:
        print("\n无'被判空场但原文本非空'的文本。")

    print("\n阶段 B 完成。")


if __name__ == "__main__":
    main()