#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""swap步数 + 重绘分辨率 对比测试（热数据）。

每个组合: 预热1次(丢弃) + 正式测1次(取热数据)。
B维度: steps=10/15/20 (分辨率固定896)
C维度: 分辨率=640/896/1024 (steps固定15)
4图×2发型=8组 × 6档 × 2次(预热+正式) = 96次
"""
import base64
import json
import os
import time
from pathlib import Path

import requests

API = "http://127.0.0.1:8187/api/v1/debug/grow-timing"
TOKEN = "dev-shared-secret-2026"
OUT = Path("/home/ubuntu/hair/benchmark_out/bench2")
OUT.mkdir(parents=True, exist_ok=True)

IMGS = [
    ("asdf", "/home/ubuntu/hair/image/asdf.jpg"),
    ("qwer", "/home/ubuntu/hair/image/qwer.jpg"),
    ("girl2", "/home/ubuntu/hair/image/girl_img/girl2.jpg"),
    ("girl5", "/home/ubuntu/hair/image/girl_img/girl5.jpg"),
]
HAIRSTYLES = [(5, "wave", "波浪"), (3, "heart", "心形")]

# B维度: swap步数对比 (分辨率固定896)
B_STEPS = [10, 15, 20]
# C维度: 重绘分辨率对比 (steps固定15)
C_RES = [640, 896, 1024]


def call(img_path, hair_num, webui_steps=None, redraw_max_side=None, save_grown=None):
    """调调试接口。返回 dict。save_grown 非None时把结果图存到该路径。"""
    data = {"hair_style": str(hair_num)}
    if webui_steps is not None:
        data["webui_steps"] = str(webui_steps)
    if redraw_max_side is not None:
        data["redraw_max_side"] = str(redraw_max_side)
    t0 = time.perf_counter()
    try:
        with open(img_path, "rb") as f:
            r = requests.post(API, headers={"X-Internal-Token": TOKEN},
                files={"image_file": (os.path.basename(img_path), f, "image/jpeg")},
                data=data, timeout=300)
        wall = time.perf_counter() - t0
        j = r.json()
        if j.get("code") != 0:
            return {"ok": False, "error": j.get("message", "")[:100], "wall": wall}
        d = j["data"]
        hs = d["per_hairstyle"][0]
        if save_grown and hs.get("grown_b64"):
            b = hs["grown_b64"].split(",")[1] if "," in hs["grown_b64"] else hs["grown_b64"]
            with open(save_grown, "wb") as gf:
                gf.write(base64.b64decode(b))
        return {
            "ok": hs.get("ok", False), "wall": wall,
            "total_ms": d["total_ms"], "ctx_ms": d["extract_context_ms"],
            "mask_ms": hs.get("mask_ms"), "swap_ms": hs.get("swap_ms"),
            "blend_ms": hs.get("blend_ms"), "comfy_ms": hs.get("comfyui_redraw_ms"),
            "error": hs.get("error"),
        }
    except Exception as e:
        return {"ok": False, "error": str(e)[:100], "wall": time.perf_counter() - t0}


def main():
    results = {"B_steps": [], "C_res": []}
    total_calls = len(IMGS) * len(HAIRSTYLES) * (len(B_STEPS) + len(C_RES)) * 2
    idx = 0

    # ===== B维度: swap步数对比 (分辨率固定896) =====
    print("\n===== B维度: swap步数对比 (分辨率=896) =====", flush=True)
    for steps in B_STEPS:
        print(f"\n--- steps={steps} ---", flush=True)
        for ilabel, ipath in IMGS:
            for hnum, hkey, hname in HAIRSTYLES:
                # 预热(丢弃)
                idx += 1
                print(f"[{idx}/{total_calls}] 预热 {ilabel}|{hname}|steps={steps}", flush=True)
                call(ipath, hnum, webui_steps=steps, redraw_max_side=896)
                # 正式(热数据)
                idx += 1
                save = OUT / f"B_steps{steps}_{ilabel}_{hkey}.jpg"
                print(f"[{idx}/{total_calls}] 正式 {ilabel}|{hname}|steps={steps}", flush=True)
                r = call(ipath, hnum, webui_steps=steps, redraw_max_side=896, save_grown=save)
                r["steps"] = steps; r["img"] = ilabel; r["hair"] = hkey; r["hair_name"] = hname
                r["grown_path"] = str(save) if r.get("ok") else None
                print(f"    -> total={r.get('total_ms')}ms swap={r.get('swap_ms')}ms comfy={r.get('comfy_ms')}ms ok={r.get('ok')}", flush=True)
                results["B_steps"].append(r)

    # ===== C维度: 重绘分辨率对比 (steps固定15) =====
    print("\n===== C维度: 重绘分辨率对比 (steps=15) =====", flush=True)
    for res in C_RES:
        print(f"\n--- res={res} ---", flush=True)
        for ilabel, ipath in IMGS:
            for hnum, hkey, hname in HAIRSTYLES:
                idx += 1
                print(f"[{idx}/{total_calls}] 预热 {ilabel}|{hname}|res={res}", flush=True)
                call(ipath, hnum, webui_steps=15, redraw_max_side=res)
                idx += 1
                save = OUT / f"C_res{res}_{ilabel}_{hkey}.jpg"
                print(f"[{idx}/{total_calls}] 正式 {ilabel}|{hname}|res={res}", flush=True)
                r = call(ipath, hnum, webui_steps=15, redraw_max_side=res, save_grown=save)
                r["res"] = res; r["img"] = ilabel; r["hair"] = hkey; r["hair_name"] = hname
                r["grown_path"] = str(save) if r.get("ok") else None
                print(f"    -> total={r.get('total_ms')}ms swap={r.get('swap_ms')}ms comfy={r.get('comfy_ms')}ms ok={r.get('ok')}", flush=True)
                results["C_res"].append(r)

    with open(OUT / "results.json", "w", encoding="utf-8") as f:
        json.dump(results, f, ensure_ascii=False, indent=2)
    ok = sum(1 for r in results["B_steps"] + results["C_res"] if r.get("ok"))
    print(f"\n✓ 完成: {ok}/{len(results['B_steps'])+len(results['C_res'])} 成功 -> {OUT/'results.json'}", flush=True)


if __name__ == "__main__":
    main()