# -*- coding: utf-8 -*-
"""
C3: B稿v2 ⊕ ASR 交叉复审 → 融合B稿(743行) + fusion_review.csv
=============================================================
职责:逐行复审 B稿(屏幕字幕OCR),以 ASR(口语转写)为上下文参考,
只做纠错,严禁改行数/时间戳。
"""

import json
import re
import sys
from pathlib import Path
from typing import List, Dict, Optional

from .llm import chat

# --------------------------------------------------------------------------
# 常量
# --------------------------------------------------------------------------

CHANGE_TYPE_ENUM = frozenset(
    {
        "unchanged",
        "minor_edit",
        "term_normalize",
        "rewrite_large",
        "segment_delete",
        "segment_add",
        "editor_typo",
    }
)

SYSTEM_PROMPT = """你是《军事科技》专题片文稿校审员。给你 B稿(屏幕字幕OCR,逐行碎句,带时间戳) 和对应的 ASR(口语转写)。
你的任务:逐行复审 B稿,只做纠错,绝不合并行、不拆行、不增删行、不改时间戳。
权威优先级:
- 屏幕术语/型号/番号(箭-3/萨德/见证者-136等): B稿为准(屏幕实打的字)
- B稿明显是OCR错字而ASR是对的: 用ASR覆盖
- ⚠️ 专有名词铁律：厂名/型号/番号/国名/人名/机构名等专名，遇B稿与ASR同音异写(如斯泰尔vs斯太尔、美以vs美伊)，一律以B稿/A稿书面写法为准，零容忍采ASR。ASR是口语转写，同音字极多，专名绝不信ASR。
- 同音事实错(如"美以"vs"美伊"): 以书面规范为准,存疑进review
- 一两个字的等价差异(的/地、啊等语气): 算 unchanged,不要改
每行输出: line_no, final_text(纠错后,默认等于B原文), change_type(7选1), confidence(0~1), reason(简短,unchanged时留空)
只返回JSON数组,不要任何解释文字。
change_type枚举: unchanged/minor_edit/term_normalize/rewrite_large/segment_delete/segment_add/editor_typo"""


# --------------------------------------------------------------------------
# 1. 解析带时间戳的行
# --------------------------------------------------------------------------

def parse_timed_lines(path) -> List[dict]:
    r"""
    解析 "[XmYs] 文本" → [{"idx":int, "ts_raw":"0m8s", "ts_sec":8, "text":"导弹呼啸而过"}]
    正则: ^\[(\d+)m(\d+)s\]\s*(.*)$ ; ts_sec = m*60+s
    解析失败的行要抛异常并打印行号,不许静默跳过
    """
    p = Path(path)
    if not p.exists():
        raise FileNotFoundError(f"文件不存在: {path}")

    pattern = re.compile(r"^\[(\d+)m(\d+)s\]\s*(.*)$")
    lines_raw = p.read_text(encoding="utf-8").splitlines()
    result = []

    for idx, line in enumerate(lines_raw, start=1):
        line = line.strip()
        if not line:
            continue  # 跳过空行
        m = pattern.match(line)
        if not m:
            raise ValueError(
                f"行 {idx} 解析失败,无法匹配时间戳格式: {repr(line[:120])}\n"
                f"文件: {path}"
            )
        minutes = int(m.group(1))
        seconds = int(m.group(2))
        ts_raw = f"{minutes}m{seconds}s"
        ts_sec = minutes * 60 + seconds
        text = m.group(3).strip()
        result.append(
            {
                "idx": idx,
                "ts_raw": ts_raw,
                "ts_sec": ts_sec,
                "text": text,
            }
        )

    return result


# --------------------------------------------------------------------------
# 2. 对齐 ASR 上下文
# --------------------------------------------------------------------------

def align_asr_context(b_lines: List[dict], asr_lines: List[dict]) -> List[str]:
    """
    为每个 B 行找时间窗内的 ASR 上下文(用于喂 LLM)
    规则: 取 ts_sec 落在 [b_ts-3, b_next_ts+3] 区间的 ASR 句拼接;
    边界用前后各扩 1 句兜底。返回与 b_lines 等长的 context 列表
    """
    n = len(b_lines)
    contexts = []

    # 预计算 B 行的时间窗: [b[i].ts_sec - 3, b[i+1].ts_sec + 3]
    # 最后一行用 b[i].ts_sec + 10 作为上界
    windows = []
    for i, bl in enumerate(b_lines):
        lo = bl["ts_sec"] - 3
        if i + 1 < n:
            hi = b_lines[i + 1]["ts_sec"] + 3
        else:
            hi = bl["ts_sec"] + 10
        windows.append((lo, hi))

    asr_count = len(asr_lines)

    for i, (lo, hi) in enumerate(windows):
        # 找到落在窗口内的 ASR 句索引
        hit_indices = []
        for j, al in enumerate(asr_lines):
            if lo <= al["ts_sec"] <= hi:
                hit_indices.append(j)

        if not hit_indices:
            # 无命中:取距离最近的 1 句
            best_j = None
            best_dist = float("inf")
            mid_ts = (lo + hi) / 2
            for j, al in enumerate(asr_lines):
                dist = abs(al["ts_sec"] - mid_ts)
                if dist < best_dist:
                    best_dist = dist
                    best_j = j
            if best_j is not None:
                start_j = max(0, best_j - 1)
                end_j = min(asr_count - 1, best_j + 1)
            else:
                start_j = 0
                end_j = 0
        else:
            # 命中句的范围 + 前后各扩 1
            start_j = max(0, hit_indices[0] - 1)
            end_j = min(asr_count - 1, hit_indices[-1] + 1)

        # 拼接 [start_j, end_j] 的 ASR 文本
        selected = asr_lines[start_j : end_j + 1]
        context = " ".join(s["text"] for s in selected)
        contexts.append(context)

    assert len(contexts) == n, f"context 列表长度 {len(contexts)} != B 稿行数 {n}"
    return contexts


# --------------------------------------------------------------------------
# 3. 构造 Prompt
# --------------------------------------------------------------------------

def build_prompt(batch_b: List[dict], batch_ctx: List[str]) -> List[dict]:
    """
    构造 messages,见下方"四、Prompt 模板"
    """
    assert len(batch_b) == len(batch_ctx), (
        f"batch_b({len(batch_b)}) 与 batch_ctx({len(batch_ctx)}) 长度不一致"
    )

    user_lines = []
    for bl, ctx in zip(batch_b, batch_ctx):
        line_no = bl["idx"]
        b_text = bl["text"]
        asr_text = ctx if ctx else "(无ASR上下文)"
        user_lines.append(
            f"[行{line_no}] B稿: \"{b_text}\"  ASR上下文: \"{asr_text}\""
        )

    user_content = "\n".join(user_lines)

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_content},
    ]
    return messages


# --------------------------------------------------------------------------
# 4. 单批复审
# --------------------------------------------------------------------------

def review_batch(
    batch_b: List[dict], batch_ctx: List[str], no_ai: bool = False
) -> List[dict]:
    """
    no_ai=True: 直接回填 unchanged(final_text=b原文, change_type="unchanged", confidence=1.0)
    no_ai=False: 调 llm.chat(messages, thinking=False, max_tokens=4000, temperature=0.0)
      解析返回 JSON 数组; 每元素 {line_no, final_text, change_type, confidence, reason}
    返回标准化记录列表
    """
    if no_ai:
        records = []
        for bl in batch_b:
            records.append(
                {
                    "line_no": bl["idx"],
                    "final_text": bl["text"],
                    "change_type": "unchanged",
                    "confidence": 1.0,
                    "reason": "",
                }
            )
        return records

    # ---- AI 路径 ----
    messages = build_prompt(batch_b, batch_ctx)

    try:
        raw_response = chat(
            messages,
            thinking=False,
            max_tokens=4000,
            temperature=0.0,
        )
    except Exception as e:
        print(
            f"[fusion_review] LLM 调用失败,回退为 unchanged 批次: {e}",
            file=sys.stderr,
        )
        # 回退:全部 unchanged
        records = []
        for bl in batch_b:
            records.append(
                {
                    "line_no": bl["idx"],
                    "final_text": bl["text"],
                    "change_type": "unchanged",
                    "confidence": 1.0,
                    "reason": f"LLM调用失败回退: {str(e)[:80]}",
                }
            )
        return records

    # 解析 JSON
    parsed = _parse_llm_json_response(raw_response, len(batch_b))

    # 标准化并校验
    records = []
    for item in parsed:
        line_no = item.get("line_no")
        final_text = item.get("final_text", "")
        change_type = item.get("change_type", "unchanged")
        confidence = item.get("confidence", 1.0)
        reason = item.get("reason", "")

        # 校验 change_type
        if change_type not in CHANGE_TYPE_ENUM:
            original_ct = change_type
            print(
                f"[fusion_review] 行 {line_no} 非法 change_type='{original_ct}', 强制改为 unchanged",
                file=sys.stderr,
            )
            change_type = "unchanged"
            final_text = ""  # 下面会回填
            reason = f"LLM返回非法change_type({original_ct}),回退unchanged"

        # 如果 change_type 被改为 unchanged 但 final_text 为空,回填 B 原文
        if change_type == "unchanged" and not final_text:
            # 从 batch_b 找回原文
            for bl in batch_b:
                if bl["idx"] == line_no:
                    final_text = bl["text"]
                    break

        records.append(
            {
                "line_no": line_no,
                "final_text": final_text,
                "change_type": change_type,
                "confidence": float(confidence),
                "reason": reason or "",
            }
        )

    return records


def _parse_llm_json_response(raw: str, expected_len: int) -> List[dict]:
    """解析 LLM 返回的 JSON,处理 markdown code fences 等常见包装。"""
    text = raw.strip()

    # 去掉可能的 markdown code fences
    if text.startswith("```"):
        lines = text.splitlines()
        # 去掉第一行 ```json 或 ```
        if lines and lines[0].startswith("```"):
            lines = lines[1:]
        # 去掉最后一行 ```
        if lines and lines[-1].strip() == "```":
            lines = lines[:-1]
        text = "\n".join(lines).strip()

    try:
        result = json.loads(text)
    except json.JSONDecodeError as e:
        raise ValueError(
            f"LLM 返回 JSON 解析失败: {e}\n"
            f"原始响应前 500 字符: {raw[:500]}"
        )

    if not isinstance(result, list):
        raise ValueError(
            f"LLM 返回不是 JSON 数组,类型为 {type(result).__name__}"
        )

    if len(result) != expected_len:
        raise ValueError(
            f"LLM 返回 {len(result)} 条记录,期望 {expected_len} 条。"
            f"该批次将回退为 unchanged 并重新请求。"
        )

    return result


# --------------------------------------------------------------------------
# 5. 主流程
# --------------------------------------------------------------------------

def run_fusion(
    episode_id: str,
    output_dir: str,
    no_ai: bool = False,
    batch_size: int = 35,
) -> dict:
    """
    主流程:
    1. 读 output_dir/B稿_v2.txt → b_lines（断言行数>0）
    2. 读 output_dir/asr_v2_timed.txt → asr_lines
    3. align_asr_context 生成等长 context
    4. 按 batch_size 分块；每块结果落缓存,已存在则复用（断点续跑）
    5. 逐块 review_batch,汇总所有记录
    6. 硬校验（任一不过就 raise,不写出文件）
    7. 写 output_dir/融合B稿.txt
    8. 写 output_dir/fusion_review.csv
    9. 返回统计 dict
    """
    out_dir = Path(output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    b_path = out_dir / "B稿_v2.txt"
    asr_path = out_dir / "asr_v2_timed.txt"

    if not b_path.exists():
        raise FileNotFoundError(f"B稿_v2.txt 不存在: {b_path}")
    if not asr_path.exists():
        raise FileNotFoundError(f"asr_v2_timed.txt 不存在: {asr_path}")

    # Step 1: 解析 B 稿
    b_lines = parse_timed_lines(b_path)
    assert len(b_lines) > 0, f"B稿_v2.txt 解析后为空: {b_path}"

    # Step 2: 解析 ASR
    asr_lines = parse_timed_lines(asr_path)

    # Step 3: 对齐 ASR 上下文
    contexts = align_asr_context(b_lines, asr_lines)
    assert len(contexts) == len(b_lines), (
        f"context 长度 {len(contexts)} != B 稿行数 {len(b_lines)}"
    )

    # Step 4: 分块 + 缓存
    cache_dir = out_dir / ".c3_cache"
    cache_dir.mkdir(parents=True, exist_ok=True)

    all_records = []
    total_batches = (len(b_lines) + batch_size - 1) // batch_size

    for batch_idx in range(total_batches):
        start = batch_idx * batch_size
        end = min(start + batch_size, len(b_lines))
        batch_b = b_lines[start:end]
        batch_ctx = contexts[start:end]

        cache_path = cache_dir / f"batch_{batch_idx}.json"

        if cache_path.exists():
            # 断点续跑:复用缓存
            try:
                cached = json.loads(cache_path.read_text(encoding="utf-8"))
                print(f"[fusion_review] 复用缓存 batch_{batch_idx} ({len(cached)} 条)")
                all_records.extend(cached)
                continue
            except Exception as e:
                print(
                    f"[fusion_review] 缓存 batch_{batch_idx} 损坏,重新计算: {e}",
                    file=sys.stderr,
                )

        print(
            f"[fusion_review] 复审 batch {batch_idx + 1}/{total_batches} "
            f"(行 {start + 1}-{end})..."
        )
        try:
            batch_records = review_batch(batch_b, batch_ctx, no_ai=no_ai)
        except Exception as e:
            print(
                f"[fusion_review] batch {batch_idx + 1} 失败,跳过缓存写入: {e}",
                file=sys.stderr,
            )
            # 不写缓存,下次重跑时重新请求该批
            continue

        # 写入缓存
        cache_path.write_text(
            json.dumps(batch_records, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        all_records.extend(batch_records)

    # Step 6: 硬校验
    _hard_validate(all_records, b_lines)

    # Step 6.5: 修正语义——final_text 等于 B 原文的行强制归为 unchanged
    _normalize_unchanged_when_no_edit(all_records, b_lines)

    # Step 7: 写 融合B稿.txt
    fused_path = out_dir / "融合B稿.txt"
    fused_lines = []
    for rec, bl in zip(all_records, b_lines):
        fused_lines.append(f"[{bl['ts_raw']}] {rec['final_text']}")
    fused_path.write_text("\n".join(fused_lines) + "\n", encoding="utf-8")

    # Step 8: 写 fusion_review.csv
    csv_path = out_dir / "fusion_review.csv"
    csv_rows = [
        "line_no,timestamp,b_original,asr_context,final_text,change_type,confidence,reason"
    ]
    for rec, bl, ctx in zip(all_records, b_lines, contexts):
        if rec["change_type"] == "unchanged" and rec["confidence"] >= 0.8:
            continue  # 只写需要 review 的行
        # CSV 转义: 字段含逗号或引号时用双引号包裹
        row_fields = [
            str(rec["line_no"]),
            bl["ts_raw"],
            bl["text"],
            ctx,
            rec["final_text"],
            rec["change_type"],
            str(rec["confidence"]),
            rec["reason"],
        ]
        csv_rows.append(_csv_row(row_fields))
    csv_path.write_text("\n".join(csv_rows) + "\n", encoding="utf-8")

    # Step 9: 统计
    stats = {
        "total_lines": len(b_lines),
        "change_counts": {},
        "review_lines": 0,
    }
    for rec in all_records:
        ct = rec["change_type"]
        stats["change_counts"][ct] = stats["change_counts"].get(ct, 0) + 1
        if ct != "unchanged" or rec["confidence"] < 0.8:
            stats["review_lines"] += 1

    print(f"[fusion_review] 融合完成: 总行数={stats['total_lines']}")
    print(f"[fusion_review] 各 change_type 计数: {stats['change_counts']}")
    print(f"[fusion_review] 进 review 行数: {stats['review_lines']}")
    print(f"[fusion_review] 融合B稿: {fused_path}")
    print(f"[fusion_review] review CSV: {csv_path}")

    stats["fused_path"] = str(fused_path)
    stats["csv_path"] = str(csv_path)

    return stats


def _hard_validate(records: List[dict], b_lines: List[dict]) -> None:
    """硬校验,任一不过就 raise ValueError,不写出文件。"""
    # 校验 1: 行数必须相等
    if len(records) != len(b_lines):
        raise ValueError(
            f"行数不一致: records={len(records)}, B稿={len(b_lines)}"
        )

    # 校验 2: 逐行时间戳不能被动
    for i, (rec, bl) in enumerate(zip(records, b_lines)):
        rec_line_no = rec.get("line_no")
        expected_line_no = bl["idx"]
        if rec_line_no != expected_line_no:
            raise ValueError(
                f"第 {i} 条记录 line_no={rec_line_no}, 期望 {expected_line_no}"
            )

    # 校验 3: change_type 枚举
    for rec in records:
        ct = rec.get("change_type", "")
        if ct not in CHANGE_TYPE_ENUM:
            raise ValueError(
                f"行 {rec.get('line_no')}: 非法 change_type='{ct}'"
            )


def _normalize_unchanged_when_no_edit(
    records: List[dict], b_lines: List[dict]
) -> None:
    """修正语义:final_text 等于 B 原文的行,强制归为 unchanged。

    LLM 有时会把"考虑后决定保留 B 稿"标成 minor_edit,
    但实际 final_text == b_original, 这不在留痕范围内。
    """
    b_text_by_idx = {bl["idx"]: bl["text"] for bl in b_lines}
    fixed = 0
    for rec in records:
        line_no = rec.get("line_no")
        b_orig = b_text_by_idx.get(line_no)
        if b_orig is not None and rec.get("final_text") == b_orig:
            if rec.get("change_type") != "unchanged":
                rec["change_type"] = "unchanged"
                rec["confidence"] = 1.0
                rec["reason"] = ""
                fixed += 1
    if fixed:
        print(
            f"[fusion_review] 修正 {fixed} 行: final_text==B原文但change_type≠unchanged, 已强制归为 unchanged"
        )


def _csv_row(fields: List[str]) -> str:
    """将字段列表格式化为 CSV 行,处理逗号和引号转义。"""
    escaped = []
    for f in fields:
        s = str(f)
        if "," in s or '"' in s or "\n" in s:
            s = s.replace('"', '""')
            s = f'"{s}"'
        escaped.append(s)
    return ",".join(escaped)