"""Phase 3 / Checkpoint 3 验收脚本：独立验证 Stage 3 风格还原链路。

复用 Phase 2 已抽出的 normal 帧，无需再跑 Trellis（省 $0.041）。
默认会挑最新一组 view_*_el0_normal.png 作为结构参考，并让用户指定原图。

用法（在 backend 目录激活 venv 后）：
    # 1. 全 6 帧批量重绘（~6 × $0.07 ≈ $0.42）
    python scripts/test_style_restore.py <原图路径>

    # 2. 只跑一帧做快速冒烟（便宜 ~$0.07）
    python scripts/test_style_restore.py <原图路径> --single

    # 3. 手动指定 normal 帧前缀（不指定时用最新一组）
    python scripts/test_style_restore.py <原图路径> --prefix view_xxxxxxxx

成功标准：
- 退出码 0
- 打印的每一帧都返回 success=True 且 url 以 /generated/restyled_ 开头
- 打开 generated/restyled_* 图片，目视检查：风格与原图一致、视角符合 az 标注
"""

from __future__ import annotations

import argparse
import asyncio
import json
import sys
from pathlib import Path

_BACKEND_ROOT = Path(__file__).resolve().parent.parent
if str(_BACKEND_ROOT) not in sys.path:
    sys.path.insert(0, str(_BACKEND_ROOT))

from dotenv import load_dotenv  # noqa: E402

load_dotenv(_BACKEND_ROOT / ".env", override=True)

from app.services import style_restorer  # noqa: E402
from app.services.image_gen import GENERATED_DIR  # noqa: E402


def _pick_latest_normal_group() -> tuple[str, list[dict]]:
    """在 generated/ 下挑一组最新的 view_<prefix>_az*_el0_normal.png，按 az 排好序。

    返回 (prefix, frames)，frames = [{az, el, normal_png, color_png}, ...]
    """
    all_normals = sorted(
        GENERATED_DIR.glob("view_*_el0_normal.png"),
        key=lambda p: p.stat().st_mtime,
        reverse=True,
    )
    if not all_normals:
        raise FileNotFoundError(
            "generated/ 下找不到 view_*_el0_normal.png；请先跑 scripts/test_frame_extract.py"
        )
    # 最新文件的 prefix：view_<8hex>_az...
    latest_name = all_normals[0].name
    # 形如 view_63731b77_az0_el0_normal.png → prefix = view_63731b77
    parts = latest_name.split("_")
    prefix = "_".join(parts[:2])
    return prefix, _load_group_by_prefix(prefix)


def _load_group_by_prefix(prefix: str) -> list[dict]:
    normals = sorted(GENERATED_DIR.glob(f"{prefix}_az*_el0_normal.png"))
    if not normals:
        raise FileNotFoundError(f"找不到前缀为 {prefix} 的 normal 帧")
    frames = []
    for n in normals:
        # 从文件名解析 az：view_63731b77_az120_el0_normal.png → az=120
        tokens = n.stem.split("_")
        az = 0
        el = 0
        for t in tokens:
            if t.startswith("az"):
                try:
                    az = int(t[2:])
                except ValueError:
                    pass
            elif t.startswith("el"):
                try:
                    el = int(t[2:])
                except ValueError:
                    pass
        color = n.with_name(n.name.replace("_normal.png", "_color.png"))
        frames.append({
            "az": az,
            "el": el,
            "normal_png": f"/generated/{n.name}",
            "color_png": f"/generated/{color.name}" if color.exists() else None,
            "frame_idx": -1,
        })
    frames.sort(key=lambda f: f["az"])
    return frames


async def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("original", help="用户原图路径（本地相对/绝对路径均可）")
    parser.add_argument("--prefix", default=None, help="normal 帧的文件名前缀（如 view_63731b77）")
    parser.add_argument("--single", action="store_true", help="只跑第一帧做快速冒烟")
    parser.add_argument("--scale", type=float, default=None, help="覆盖 IP-Adapter 权重（0-1）")
    parser.add_argument("--cn-scale", type=float, default=None, help="覆盖 ControlNet 权重（0-1）")
    args = parser.parse_args()

    original = args.original
    original_path = Path(original)
    if not original_path.is_absolute():
        original_path = (_BACKEND_ROOT / original).resolve()
    if not original_path.exists():
        print(f"[test_style_restore] 原图不存在：{original_path}")
        return 2

    if args.prefix:
        frames = _load_group_by_prefix(args.prefix)
        prefix = args.prefix
    else:
        prefix, frames = _pick_latest_normal_group()

    if args.single:
        frames = frames[:1]

    override_params = {}
    if args.scale is not None:
        override_params["scale"] = args.scale
    if args.cn_scale is not None:
        override_params["controlnet_conditioning_scale"] = args.cn_scale

    print(f"[test_style_restore] 原图：{original_path}")
    print(f"[test_style_restore] normal 帧组 prefix={prefix}，共 {len(frames)} 帧")
    print(f"[test_style_restore] 角度列表：{[f['az'] for f in frames]}")
    if override_params:
        print(f"[test_style_restore] 覆盖参数：{override_params}")
    print(
        f"[test_style_restore] 预计耗时 ~{72}s × {len(frames)}，"
        f"成本 ~${0.07 * len(frames):.2f}（并发 3 实际会更快）"
    )

    results = await style_restorer.restore_batch(
        original_image=str(original_path),
        frames=frames,
        params=override_params or None,
        concurrency=3,
    )

    print("[test_style_restore] 重绘结果：")
    print(json.dumps(results, ensure_ascii=False, indent=2))

    ok = sum(1 for r in results if r.get("success"))
    print(f"[test_style_restore] 成功 {ok}/{len(results)} 帧")

    if ok == 0:
        print("[test_style_restore] 全军覆没，判定失败")
        return 1
    if ok < len(results):
        print("[test_style_restore] 部分帧失败，视为部分通过；请检查上方错误信息")
        return 0

    print(
        "[test_style_restore] 全部成功；请目视检查 generated/restyled_* 图片："
        "风格应贴近原图、视角按 az 顺序环绕。"
    )
    return 0


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