#!/usr/bin/env python3
"""Phase 4 中文润色：按 H2 section 切块 → 逐块润色 → 拼接 → 写盘。

用法：
    uv run python scripts/polish.py <project_slug>

架构跟 translate.py 一致：
- Python 做切块/循环/重试/断点续传
- LLM 只做"润色这一段"，单次 output token 远低于上限
- 结果落在 phase4/zh_polished_chunks/<anchor>.md，重跑只补缺
- 最终拼接写入 phase4/final_zh_polished.md（默认覆写 final_zh.md 的副本）

区别只在：
- 输入是中文（final_zh.md），输出还是中文
- 不维护术语表（翻译阶段已经固定了）
- 会附带一份 notes.jsonl 记录模型发现的异常
"""

from __future__ import annotations

import argparse
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" / "polish_system.txt"


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


def parse_polish_response(text: str) -> tuple[str, str]:
    """解析 <<<POLISHED>>>/<<<NOTES>>> 格式。返回 (polished, notes)。"""
    p_start = text.find("<<<POLISHED>>>")
    p_end = text.find("<<<END_POLISHED>>>")
    if p_start == -1 or p_end == -1 or p_end <= p_start:
        raise ValueError(f"missing <<<POLISHED>>> markers: {text[:300]}")
    polished = text[p_start + len("<<<POLISHED>>>"): p_end].strip("\r\n")

    n_start = text.find("<<<NOTES>>>")
    n_end = text.find("<<<END_NOTES>>>")
    notes = ""
    if n_start != -1 and n_end != -1 and n_end > n_start:
        notes = text[n_start + len("<<<NOTES>>>"): n_end].strip()
    return polished, notes


def build_user_prompt(block: MarkdownBlock) -> str:
    level_hint = f"H{block.level}" if block.level >= 1 else "frontmatter (无标题)"
    return (
        f"# 待润色的中文块（Markdown, {level_hint}）\n"
        "请按系统提示的规则润色下面这段中文。严格使用指定的分隔符格式输出。\n\n"
        "----- BEGIN BLOCK -----\n"
        f"{block.content}\n"
        "----- END BLOCK -----\n"
    )


def polish_block(
    client: ZenMuxClient,
    block: MarkdownBlock,
    *,
    model: str,
    system_prompt: str,
    temperature: float,
) -> tuple[str, str]:
    user = build_user_prompt(block)
    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"polish:{block.anchor}",
    )
    try:
        polished, notes = parse_polish_response(raw)
    except Exception as e:
        raise RuntimeError(
            f"bad response format for block {block.anchor}: {e}\nraw head: {raw[:300]}"
        )
    if not polished.strip():
        raise RuntimeError(f"empty polished content for block {block.anchor}")
    return polished.rstrip(), notes


def main() -> int:
    parser = argparse.ArgumentParser(description="Phase 4 中文润色（按 H2 切块循环）")
    parser.add_argument("project", help="项目 slug 或完整路径")
    parser.add_argument(
        "--source",
        default="phase4/final_zh.md",
        help="中文源（默认 phase4/final_zh.md，即 translate.py 的产物）",
    )
    parser.add_argument(
        "--output",
        default="phase4/final_zh_polished.md",
        help="润色后输出（默认 phase4/final_zh_polished.md）",
    )
    parser.add_argument("--model", default=DEFAULT_MODEL)
    parser.add_argument("--temperature", type=float, default=0.4)
    parser.add_argument(
        "--force",
        action="store_true",
        help="忽略 zh_polished_chunks 缓存，强制重润",
    )
    parser.add_argument(
        "--only",
        default=None,
        help="只润色指定 order（逗号分隔），例如 --only 2,7,18",
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=None,
        help="最多润色前 N 个未缓存的块（调试用）",
    )
    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}. 请先跑 translate.py")

    chunks_dir = project_root / "phase4" / "zh_polished_chunks"
    chunks_dir.mkdir(parents=True, exist_ok=True)
    logs_dir = project_root / "phase4" / "logs"
    log_file = logs_dir / "polish.jsonl"
    notes_file = project_root / "phase4" / "polish_notes.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)

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

    total_cn = sum(count_chinese_chars(b.content) for b in blocks)
    print(f"Source: {src_path.relative_to(project_root)}")
    print(f"Blocks: {len(blocks)}  | total Chinese chars: {total_cn:,}")
    print(f"Model: {args.model}  | temperature: {args.temperature}")
    print()

    start = time.time()
    translated_this_run = 0
    notes_records: list[dict] = []
    with ZenMuxClient(log_file=log_file) as client:
        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:
                print(f"  [ok  ] #{b.order:03d} {b.short_title}  (cached)")
                continue
            if args.limit is not None and translated_this_run >= args.limit:
                continue

            label = f"#{b.order:03d} L{b.level} {count_chinese_chars(b.content):>4}字 {b.short_title}"
            print(f"  [... ] {label}  ", end="", flush=True)
            t0 = time.time()
            try:
                polished, notes = polish_block(
                    client,
                    b,
                    model=args.model,
                    system_prompt=system_prompt,
                    temperature=args.temperature,
                )
            except (ZenMuxError, RuntimeError) as e:
                print(f"\n  [FAIL] {label}\n         {e}")
                continue
            elapsed = time.time() - t0

            chunk_path.write_text(polished + "\n", encoding="utf-8")
            cn = count_chinese_chars(polished)
            before_cn = count_chinese_chars(b.content)
            delta = cn - before_cn
            sign = "+" if delta >= 0 else ""
            translated_this_run += 1
            if notes:
                notes_records.append(
                    {"order": b.order, "anchor": b.anchor, "title": b.short_title, "notes": notes}
                )
            print(f"\r  [done] {label}  → {cn}字 ({sign}{delta}, {elapsed:4.1f}s)")

    # 汇总
    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_polished.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")

    if notes_records:
        notes_file.write_text(
            "\n".join(json.dumps(r, ensure_ascii=False) for r in notes_records) + "\n",
            encoding="utf-8",
        )
        print(f"\n发现 {len(notes_records)} 条润色笔记 → {notes_file.relative_to(project_root)}")

    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"  润色前字数: {total_cn:,}")
    print(f"  润色后字数: {cn:,}  (变化 {cn - total_cn:+d}, {(cn/total_cn - 1)*100:+.1f}%)")
    print(f"  耗时: {total_time:.1f}s")
    print(client.usage.summary())
    return 0


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