"""翻译模型横向评测执行器（真实节点代码，模型是唯一变量）。

设计原则（对齐 AGENTS.md 与本仓库测试约定）
------------------------------------------------
1. **不改生产逻辑**：评测直接调用 `nodes/llm.py` 的真实实现
   （`_system_prompt` / `translate_lines` / `invoke`：提示词、专名规则注入、
   ID 严格校验、重试、幻觉清洗、SRT 写出全部走生产路径）；被评测的模型只通过
   `params["model"]` 传入。
2. **只在 I/O 边界做手脚**（AGENTS 允许的最小 mock 面）：包装
   `urllib.request.urlopen`，用于
     a) 记录每次请求的耗时与 token 用量（生产日志只有批级耗时，这里要精确到
        单次调用，用于"每行摊薄秒数"判定）；
     b) `strip_thinking=True` 时把请求体里的 `enable_thinking` 键删掉——
        Qwen3-VL 系列不接受该参数，生产代码固定携带（实测返回 400
        code 20015）。这是**评测侧兼容**，生产代码保持原样，报告里单独标注。
3. **不并发**：llm-translate 生产实现是逐批串行的，评测同样串行（用户已确认
   本次不评估并发能力）。

用法
----
    # 用某个模型跑完整 SRT（产出 cn.srt + calls.jsonl + summary.json）
    uv run python scripts/bench_translate_models.py run \
        --model Qwen/Qwen3-14B --tag cloud-qwen3-14b

    # 本地 ollama（OpenAI 兼容端点）
    uv run python scripts/bench_translate_models.py run \
        --model qwen3:14b --tag local-qwen3-14b \
        --api-base http://localhost:11434/v1/chat/completions --api-key ""

    # 评测窗口对照（打印各模型在同一窗口的译文，供人工 review）
    uv run python scripts/bench_translate_models.py compare --tags cloud-qwen3-14b local-qwen3-14b
"""

from __future__ import annotations

import argparse
import json
import os
import time
import urllib.request
from contextlib import contextmanager
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parent.parent
# 所有评测产物落在 data/experiments（gitignored），不污染测试数据。
OUT_ROOT = PROJECT_ROOT / "data/experiments/translate_models"

# 评测输入：run_d386ccf124f7 的日语 ASR（与 run_479b411f299d 同音轨，已核对
# 前 5 分钟 16k 单声道 PCM md5 一致），1161 条 cue。
DEFAULT_INPUT = PROJECT_ROOT / "data/storage/runs/run_d386ccf124f7/steps/asr/transcript.srt"


@contextmanager
def _instrument(model: str, log_path: Path, strip_thinking: bool):
    """包装 urlopen：记录每次 LLM 请求耗时/token，并可按需剔除 enable_thinking。

    只替换 `urllib.request.urlopen` 这一个 I/O 入口，返回对象与异常语义保持
    不变；非 LLM 请求（如 ollama 其它端点）原样透传。
    """
    original = urllib.request.urlopen
    calls: list[dict] = []
    fh = log_path.open("w", encoding="utf-8")

    def patched(request, *args, **kwargs):
        body = None
        if hasattr(request, "data") and request.data:
            try:
                body = json.loads(request.data.decode("utf-8"))
            except (ValueError, UnicodeDecodeError):
                body = None
        # 评测侧兼容：Qwen3-VL 不接受 enable_thinking（生产代码固定携带）。
        had_thinking = isinstance(body, dict) and "enable_thinking" in body
        if strip_thinking and had_thinking:
            body.pop("enable_thinking")
            request.data = json.dumps(body).encode("utf-8")
        started = time.monotonic()
        record: dict = {"model": model, "had_enable_thinking": had_thinking}
        try:
            response = original(request, *args, **kwargs)
            payload = response.read()
            elapsed = time.monotonic() - started
            record["elapsed_s"] = round(elapsed, 3)
            try:
                parsed = json.loads(payload)
                usage = parsed.get("usage") or {}
                record["prompt_tokens"] = usage.get("prompt_tokens")
                record["completion_tokens"] = usage.get("completion_tokens")
                record["total_tokens"] = usage.get("total_tokens")
                record["finish_reason"] = (parsed.get("choices") or [{}])[0].get("finish_reason")
                record["ok"] = True
            except ValueError:
                record["ok"] = True
            calls.append(record)
            fh.write(json.dumps(record, ensure_ascii=False) + "\n")
            fh.flush()

            class _Replay:
                """把已读出的响应体重新包装成文件式对象，交给上层原样消费。"""

                def __init__(self, inner, data: bytes):
                    self._inner = inner
                    self._data = data

                def read(self, *a):
                    return self._data

                def __enter__(self):
                    # 生产代码用 with urlopen(...) as response，包装对象必须
                    # 支持上下文管理器协议，否则断链。
                    return self

                def __exit__(self, *a):
                    return self._inner.__exit__(*a)

                def __getattr__(self, name):
                    return getattr(self._inner, name)

            return _Replay(response, payload)
        except Exception as exc:  # noqa: BLE001 - 记录后原样抛出，行为不变
            record["elapsed_s"] = round(time.monotonic() - started, 3)
            record["ok"] = False
            record["error"] = f"{type(exc).__name__}: {exc}"
            calls.append(record)
            fh.write(json.dumps(record, ensure_ascii=False) + "\n")
            fh.flush()
            raise

    urllib.request.urlopen = patched
    try:
        yield calls
    finally:
        urllib.request.urlopen = original
        fh.close()


def run_model(args: argparse.Namespace) -> None:
    """用指定模型跑完整翻译，产出 cn.srt / calls.jsonl / summary.json。"""
    from dotenv import load_dotenv

    from wov_sdk.models import InvokeRequest
    from nodes import llm

    load_dotenv(PROJECT_ROOT / ".env")
    if args.api_base:
        os.environ["LLM_API_BASE"] = args.api_base
    # 本地 ollama 不需要密钥；显式传空串时清掉环境里的云端 key，避免误带。
    if args.api_key is not None:
        os.environ["LLM_API_KEY"] = args.api_key
    if args.timeout:
        os.environ["LLM_TIMEOUT_SECONDS"] = str(args.timeout)

    out_dir = OUT_ROOT / args.tag
    out_dir.mkdir(parents=True, exist_ok=True)
    input_path = Path(args.input)
    # --limit 仅用于冒烟验证（会截断 SRT 到前 N 条 cue）。
    if args.limit:
        from nodes.srt import parse_srt, serialize_srt

        cues = parse_srt(input_path.read_text(encoding="utf-8"))[: args.limit]
        input_path = out_dir / f"input_first{args.limit}.srt"
        input_path.write_text(serialize_srt(cues), encoding="utf-8")

    if args.warmup:
        # 预热：本地 ollama 首次请求包含模型加载（实测 qwen3:14b 约 10s），
        # 会把"每行摊薄秒数"算高。预热会丢弃结果，只让模型常驻显存。
        from nodes.srt import parse_srt, serialize_srt

        warm_input = out_dir / "warmup.srt"
        warm_input.write_text(
            serialize_srt(parse_srt(input_path.read_text(encoding="utf-8"))[:2]), encoding="utf-8"
        )
        warm_started = time.monotonic()
        llm.invoke(
            InvokeRequest(
                run_id=f"warmup_{args.tag}",
                node_instance_id="translate",
                inputs={"srt_uri": str(warm_input)},
                params={"model": args.model, "target_language": "zh-CN"},
                output_dir=str(out_dir / "warmup_steps"),
            )
        )
        print(f"[warmup] {time.monotonic() - warm_started:.1f}s")

    print(f"[run] model={args.model} tag={args.tag} input={input_path.name} "
          f"strip_thinking={args.strip_thinking} base={os.environ.get('LLM_API_BASE')}")
    started = time.monotonic()
    with _instrument(args.model, out_dir / "calls.jsonl", args.strip_thinking) as calls:
        response = llm.invoke(
            InvokeRequest(
                run_id=f"bench_{args.tag}",
                node_instance_id="translate",
                inputs={"srt_uri": str(input_path)},
                params={"model": args.model, "target_language": "zh-CN"},
                output_dir=str(out_dir / "steps"),
            )
        )
    wall = time.monotonic() - started

    summary = {
        "model": args.model,
        "tag": args.tag,
        "api_base": os.environ.get("LLM_API_BASE"),
        "strip_thinking": args.strip_thinking,
        "input": str(input_path),
        "status": response.status,
        "error": response.error,
        "wall_s": round(wall, 2),
        "calls": len(calls),
        "failed_calls": sum(1 for c in calls if not c.get("ok")),
        "total_tokens": sum(int(c.get("total_tokens") or 0) for c in calls),
    }
    if response.status == "completed":
        srt_path = Path(response.outputs["cn_srt_uri"])
        cues = sum(1 for line in srt_path.read_text(encoding="utf-8").splitlines() if "-->" in line)
        summary["cn_srt"] = str(srt_path)
        summary["cues_out"] = cues
        summary["per_cue_s"] = round(wall / cues, 3) if cues else None
        summary["cue_per_s"] = round(cues / wall, 2) if wall else None
    (out_dir / "summary.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=1), encoding="utf-8"
    )
    print(json.dumps(summary, ensure_ascii=False, indent=1))


def _window_rows(tags: list[str], limit: int, only: str | None, emit: str | None = None,
                 eval_set: str | None = None) -> None:
    """打印/写出评测窗口的多模型译文对照表（人工 review 用）。

    默认使用人工确认的评测集 scripts/data/translate_eval/eval_set.jsonl（124 个窗口，
    按时间分层抽样）；未提供时回退到候选池 windows.jsonl。

    stdout 与 markdown 文件**共用同一批文本行**（单一出口）：此前两套格式
    各自拼接，导致 review 材料与实际打印内容漂移（回归见
    tests/test_translate_model_bench.py 的对照表一致性用例）。
    """
    from nodes.srt import parse_srt

    source = Path(eval_set) if eval_set else PROJECT_ROOT / "scripts/data/translate_eval/eval_set.jsonl"
    if not source.is_file():
        source = PROJECT_ROOT / "scripts/data/translate_eval/windows.jsonl"
    windows = [json.loads(line) for line in source.read_text(encoding="utf-8").splitlines()]
    ja_cues = parse_srt(DEFAULT_INPUT.read_text(encoding="utf-8"))
    # 每个 tag 读 cn.srt，按"时间戳 -> 译文"建索引（翻译不改变时间戳；被幻觉
    # 清洗删掉的 cue 缺失属预期，用 None 表示）。
    outputs: dict[str, dict[str, str]] = {}
    for tag in tags:
        path = OUT_ROOT / tag / "steps" / "cn.srt"
        if not path.is_file():
            print(f"!! 缺少 {path}")
            continue
        outputs[tag] = {
            f"{c.start}-->{c.end}": c.text for c in parse_srt(path.read_text(encoding="utf-8"))
        }
    selected = windows if only is None else [w for w in windows if only in w["zh_start"]]
    lines_out: list[str] = []
    for window in selected[: limit or None]:
        ja = window["ja_lines"]
        keys = [f"{ja_cues[i - 1].start}-->{ja_cues[i - 1].end}" for i in window["ja_ids"]]
        header = f"### {window.get('id', '?')}. {window['zh_start']} [{len(ja)}条]"
        lines_out.append(header)
        lines_out.append("- JA: " + " ｜ ".join(ja))
        lines_out.append(f"- 基准中文(OCR): {window['zh_ref']}")
        for tag, index in outputs.items():
            texts = [index.get(k) or "<已删除>" for k in keys]
            lines_out.append(f"- {tag}: " + " ｜ ".join(texts))
        lines_out.append("")
    text = "\n".join(lines_out)
    print(text)
    if emit:
        Path(emit).write_text(text, encoding="utf-8")
        print(f"\n对照表已写入 {emit}")


def main() -> None:
    parser = argparse.ArgumentParser(description="翻译模型横向评测")
    sub = parser.add_subparsers(dest="command", required=True)

    run = sub.add_parser("run", help="用某模型跑完整翻译")
    run.add_argument("--model", required=True)
    run.add_argument("--tag", required=True, help="产物目录名（data/experiments/translate_models/<tag>）")
    run.add_argument("--input", default=str(DEFAULT_INPUT))
    run.add_argument("--api-base", default=None, help="覆盖 LLM_API_BASE（本地 ollama 用）")
    run.add_argument("--api-key", default=None, help="覆盖 LLM_API_KEY（本地传空串）")
    run.add_argument("--timeout", type=float, default=None, help="LLM_TIMEOUT_SECONDS")
    run.add_argument("--strip-thinking", action="store_true", help="请求体剔除 enable_thinking（VL 模型）")
    run.add_argument("--warmup", action="store_true", help="先跑 2 条预热（本地模型加载耗时不计入）")
    run.add_argument("--limit", type=int, default=0, help="只跑前 N 条 cue（冒烟）")
    run.set_defaults(func=run_model)

    cmp_parser = sub.add_parser("compare", help="打印窗口级多模型译文对照")
    cmp_parser.add_argument("--tags", nargs="+", required=True)
    cmp_parser.add_argument("--limit", type=int, default=0)
    cmp_parser.add_argument("--only", default=None, help="只看时间戳包含该字符串的窗口")
    cmp_parser.add_argument("--emit", default=None, help="把对照表写入 markdown 文件")
    cmp_parser.add_argument("--eval-set", default=None, help="改用指定评测集 jsonl")
    cmp_parser.set_defaults(func=lambda a: _window_rows(a.tags, a.limit, a.only, a.emit, a.eval_set))

    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()