#!/usr/bin/env python3
"""Phase 4 英译中：章节级切块 → 逐块翻译 → 拼接 → 写盘。

用法：
    uv run python scripts/translate.py <project_slug>
    # 或：
    uv run python scripts/translate.py projects/dual-target-rnai-pipeline-2026

断点续传：每块翻译完立即写入 `phase4/zh_chunks/<anchor>.md`，术语表 patch 在本轮结束后统一合并。
重跑时已存在的块直接跳过，只译缺的。

设计要点：
1. 切块按 H2 粒度，单块一般 <600 英文词，单次 API 调用远低于 Sonnet output token 上限
2. 并发翻译使用稳定术语表快照，译完后统一合并 glossary patch，避免多线程写冲突
3. 失败不会污染最终产物：块级文件独立，可重跑；汇总步骤独立
4. 日志完整：每次 API 调用写 `phase4/logs/translate.jsonl`
"""

from __future__ import annotations

import argparse
import concurrent.futures
import json
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from scripts.lib.markdown_chunker import (
    MarkdownBlock,
    count_chinese_chars,
    split_by_headers,
)
from scripts.lib.zenmux_client import ZenMuxClient, ZenMuxError, load_secrets

DEFAULT_MODEL = "anthropic/claude-sonnet-4.6"
MODEL_MAX_TOKENS = {
    "anthropic/claude-sonnet-4.6": 32000,
    "anthropic/claude-sonnet-4.5": 32000,
    "anthropic/claude-opus-4.7": 32000,
    "anthropic/claude-opus-4.6": 32000,
    "anthropic/claude-haiku-4.5": 16000,
}
PROMPT_FILE = Path(__file__).parent / "prompts" / "translate_system.txt"


def resolve_project(arg: str) -> Path:
    p = Path(arg)
    if p.is_dir():
        return p
    here = Path.cwd()
    cand = here / "projects" / arg
    if cand.is_dir():
        return cand
    raise SystemExit(f"project not found: {arg}")


def parse_delimited_response(text: str) -> tuple[str, dict[str, str]]:
    """解析自定义分隔符格式的响应。

    期望结构：
        <<<TRANSLATION>>>
        ...markdown...
        <<<END_TRANSLATION>>>
        <<<GLOSSARY_PATCH>>>
        English || 中文
        ...
        <<<END_GLOSSARY_PATCH>>>

    Returns:
        (translation, glossary_patch)
    """
    t_start = text.find("<<<TRANSLATION>>>")
    t_end = text.find("<<<END_TRANSLATION>>>")
    if t_start == -1 or t_end == -1 or t_end <= t_start:
        raise ValueError(
            f"missing <<<TRANSLATION>>> markers in response: {text[:300]}"
        )
    translation = text[t_start + len("<<<TRANSLATION>>>"): t_end].strip("\r\n")

    g_start = text.find("<<<GLOSSARY_PATCH>>>")
    g_end = text.find("<<<END_GLOSSARY_PATCH>>>")
    patch: dict[str, str] = {}
    if g_start != -1 and g_end != -1 and g_end > g_start:
        body = text[g_start + len("<<<GLOSSARY_PATCH>>>"): g_end]
        for line in body.splitlines():
            line = line.strip()
            if not line or "||" not in line:
                continue
            en, _, zh = line.partition("||")
            en, zh = en.strip(), zh.strip()
            if en and zh:
                patch[en] = zh
    return translation, patch


def load_glossary(path: Path) -> dict[str, str]:
    if not path.exists():
        return {}
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return {}


def save_glossary(path: Path, glossary: dict[str, str]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(
        json.dumps(glossary, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )


def _glossary_value(value) -> str:
    if isinstance(value, dict):
        return str(value.get("zh") or value.get("中文") or "")
    return str(value)


def build_user_prompt(block: MarkdownBlock, glossary: dict[str, str]) -> str:
    glossary_lines = (
        f"{en} || {_glossary_value(zh)}"
        for en, zh in sorted(glossary.items())
        if _glossary_value(zh)
    )
    glossary_hint = "\n".join(glossary_lines) or "(none yet)"
    level_hint = (
        f"H{block.level}" if block.level >= 1 else "frontmatter (no heading)"
    )
    return (
        "# GLOSSARY (English || Chinese, already used earlier in the document):\n"
        f"{glossary_hint}\n\n"
        f"# BLOCK TO TRANSLATE (Markdown, {level_hint}):\n"
        "Translate the block between the markers below into Chinese, "
        "following all rules in the system prompt. Output using the exact "
        "delimiter format specified.\n\n"
        "----- BEGIN BLOCK -----\n"
        f"{block.content}\n"
        "----- END BLOCK -----\n"
    )


def translate_block(
    client: ZenMuxClient,
    block: MarkdownBlock,
    glossary: dict[str, str],
    *,
    model: str,
    system_prompt: str,
    temperature: float,
) -> tuple[str, dict[str, str]]:
    user = build_user_prompt(block, glossary)
    max_tok = MODEL_MAX_TOKENS.get(model, 16000)
    raw = client.chat_complete(
        model=model,
        system=system_prompt,
        user=user,
        temperature=temperature,
        max_tokens=max_tok,
        tag=f"translate:{block.anchor}",
    )
    try:
        translation, patch = parse_delimited_response(raw)
    except Exception as e:
        raise RuntimeError(
            f"bad response format for block {block.anchor}: {e}\nraw head: {raw[:300]}"
        )
    if not translation.strip():
        raise RuntimeError(f"empty translation for block {block.anchor}")
    return translation.rstrip(), patch


def main() -> int:
    parser = argparse.ArgumentParser(description="Phase 4 英译中（章节级切块并行翻译）")
    parser.add_argument("project", help="项目 slug 或完整路径")
    parser.add_argument(
        "--source",
        default="phase4/final_en.md",
        help="英文源文件（相对项目根，默认 phase4/final_en.md）",
    )
    parser.add_argument(
        "--output",
        default="phase4/final_zh.md",
        help="中文输出（默认 phase4/final_zh.md）",
    )
    parser.add_argument("--model", default=DEFAULT_MODEL, help="翻译模型")
    parser.add_argument("--temperature", type=float, default=0.3)
    parser.add_argument(
        "--force",
        action="store_true",
        help="忽略已有 zh_chunks 缓存，强制重翻",
    )
    parser.add_argument(
        "--only",
        default=None,
        help="只翻译指定 order（逗号分隔的整数），其余跳过（调试用），例如 --only 0,1,7",
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=None,
        help="最多翻译前 N 个未缓存的块（调试用）",
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=4,
        help="并发翻译 worker 数（默认 4；设为 1 回退串行）",
    )
    args = parser.parse_args()

    load_secrets()
    project_root = resolve_project(args.project)
    src_path = project_root / args.source
    out_path = project_root / args.output
    if not src_path.exists():
        raise SystemExit(f"source not found: {src_path}")

    chunks_dir = project_root / "phase4" / "zh_chunks"
    chunks_dir.mkdir(parents=True, exist_ok=True)
    glossary_path = project_root / "phase4" / "glossary.json"
    logs_dir = project_root / "phase4" / "logs"
    log_file = logs_dir / "translate.jsonl"

    system_prompt = PROMPT_FILE.read_text(encoding="utf-8")
    text = src_path.read_text(encoding="utf-8")
    blocks = split_by_headers(text, max_level=2)
    glossary = load_glossary(glossary_path)

    only_orders: set[int] | None = None
    if args.only:
        only_orders = {int(a.strip()) for a in args.only.split(",") if a.strip()}

    total_en_words = sum(b.word_count for b in blocks)
    print(f"Source: {src_path.relative_to(project_root)}")
    print(f"Blocks: {len(blocks)}  | total English words: {total_en_words:,}")
    print(f"Glossary loaded: {len(glossary)} terms")
    workers = max(1, args.workers)
    print(f"Model: {args.model}  | temperature: {args.temperature}  | workers: {workers}")
    print()

    start = time.time()
    todo: list[MarkdownBlock] = []
    cached_count = 0
    for b in blocks:
        chunk_path = chunks_dir / f"{b.order:03d}-{b.anchor}.md"
        if only_orders is not None and b.order not in only_orders:
            continue
        if chunk_path.exists() and not args.force:
            cached_count += 1
            print(f"  [ok  ] #{b.order:03d} {b.short_title}  (cached)")
            continue
        if args.limit is not None and len(todo) >= args.limit:
            continue
        todo.append(b)

    if todo:
        print(f"To translate this run: {len(todo)} blocks  | cached: {cached_count}")
        print()

    successful_patches: list[dict[str, str]] = []
    with ZenMuxClient(log_file=log_file) as client:
        def run_one(b: MarkdownBlock) -> tuple[MarkdownBlock, str, dict[str, str], float]:
            chunk_path = chunks_dir / f"{b.order:03d}-{b.anchor}.md"
            t0 = time.time()
            translation, patch = translate_block(
                client,
                b,
                glossary,
                model=args.model,
                system_prompt=system_prompt,
                temperature=args.temperature,
            )
            chunk_path.write_text(translation + "\n", encoding="utf-8")
            return b, translation, patch, time.time() - t0

        if workers == 1:
            for b in todo:
                label = f"#{b.order:03d} L{b.level} {b.word_count:>4}w  {b.short_title}"
                print(f"  [... ] {label}  ", end="", flush=True)
                try:
                    _block, translation, patch, elapsed = run_one(b)
                except (ZenMuxError, RuntimeError) as e:
                    print(f"\n  [FAIL] {label}\n         {e}")
                    continue
                successful_patches.append(patch)
                cn = count_chinese_chars(translation)
                print(f"\r  [done] {label}  → {cn:>4}字  ({elapsed:4.1f}s, +{len(patch)} terms)")
        else:
            with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as pool:
                futures = {pool.submit(run_one, b): b for b in todo}
                for fut in concurrent.futures.as_completed(futures):
                    b = futures[fut]
                    label = f"#{b.order:03d} L{b.level} {b.word_count:>4}w  {b.short_title}"
                    try:
                        _block, translation, patch, elapsed = fut.result()
                    except (ZenMuxError, RuntimeError) as e:
                        print(f"  [FAIL] {label}\n         {e}")
                        continue
                    successful_patches.append(patch)
                    cn = count_chinese_chars(translation)
                    print(f"  [done] {label}  → {cn:>4}字  ({elapsed:4.1f}s, +{len(patch)} terms)")

        glossary_conflicts: list[tuple[str, str, str]] = []
        for patch in successful_patches:
            for k, v in patch.items():
                if k not in glossary:
                    glossary[k] = v
                elif _glossary_value(glossary[k]) != v:
                    glossary_conflicts.append((k, _glossary_value(glossary[k]), v))
        if successful_patches:
            save_glossary(glossary_path, glossary)
        if glossary_conflicts:
            print(f"  [WARN] glossary patch 冲突 {len(glossary_conflicts)} 条，保留既有译法")

    # 汇总：按 order 拼接所有 chunk
    merged: list[str] = []
    missing: list[str] = []
    for b in blocks:
        chunk_path = chunks_dir / f"{b.order:03d}-{b.anchor}.md"
        if not chunk_path.exists():
            missing.append(f"#{b.order:03d} {b.short_title}")
            continue
        merged.append(chunk_path.read_text(encoding="utf-8").rstrip())

    partial = only_orders is not None or args.limit is not None
    if missing:
        print(f"\n⚠ 缺失 {len(missing)} 块：")
        for m in missing[:10]:
            print(f"  - {m}")
        if len(missing) > 10:
            print(f"  ... 还有 {len(missing) - 10} 块")
        if partial:
            print("（partial 模式：--only / --limit 生效，未生成 final_zh.md）")
        else:
            print("重新运行本脚本即可补译（已译的会跳过）。")
        print(client.usage.summary())
        return 1

    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text("\n\n".join(merged) + "\n", encoding="utf-8")

    # 统计
    final_text = out_path.read_text(encoding="utf-8")
    cn = count_chinese_chars(final_text)
    total_time = time.time() - start
    print()
    print(f"✓ 输出：{out_path.relative_to(project_root)}")
    print(f"  中文字数: {cn:,}")
    print(f"  英文词数源: {total_en_words:,}   膨胀率: {cn / max(total_en_words,1):.2f}×")
    print(f"  耗时: {total_time:.1f}s")
    print(f"  术语表: {len(glossary)} 条 → {glossary_path.relative_to(project_root)}")
    print(client.usage.summary())
    return 0


if __name__ == "__main__":
    sys.exit(main())