#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""发型对比矩阵测试：3图×5发型=15行，每行10张图(4b@896×1 + 9b三模型×三分辨率×9)。
按模型分组跑（减少模型切换次数、降低OOM风险），结果重组为15行存JSON+生成报告。
"""
import base64
import json
import os
import subprocess
import time
from collections import defaultdict
from pathlib import Path

import requests

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

IMGS = [
    ("asdf", "/home/ubuntu/hair/image/asdf.jpg"),
    ("qwer", "/home/ubuntu/hair/image/qwer.jpg"),
    ("girl5", "/home/ubuntu/hair/image/girl_img/girl5.jpg"),
]
HAIRSTYLES = [
    (1, "ellipse", "椭圆"), (2, "flower", "花瓣"), (3, "heart", "心形"),
    (4, "straight", "直线"), (5, "wave", "波浪"),
]

# 按模型分组：每个模型对应其要跑的(分辨率,列标题)
MODEL_GROUPS = [
    ("flux-2-klein-4b-fp8.safetensors", [("896", "4B@896")]),
    ("flux2.0/flux-2-klein-9b-fp8.safetensors",
     [("0", "9B-fp8@原图"), ("896", "9B-fp8@896"), ("640", "9B-fp8@640")]),
    ("flux-2-klein-9b-Q5_K_M.gguf",
     [("0", "9B-Q5@原图"), ("896", "9B-Q5@896"), ("640", "9B-Q5@640")]),
    ("flux-2-klein-9b-Q4_K_M.gguf",
     [("0", "9B-Q4@原图"), ("896", "9B-Q4@896"), ("640", "9B-Q4@640")]),
]
# 列顺序（4b在前，然后9b三模型）
COLUMN_TITLES = ["4B@896", "9B-fp8@原图", "9B-fp8@896", "9B-fp8@640",
                 "9B-Q5@原图", "9B-Q5@896", "9B-Q5@640",
                 "9B-Q4@原图", "9B-Q4@896", "9B-Q4@640"]


def gpu_used():
    try:
        out = subprocess.check_output(
            ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"], timeout=10)
        return int(out.decode().strip())
    except Exception:
        return 0


def call(img_path, hair_num, model_file, res_val):
    fd = {"gender": "female", "hair_style": str(hair_num), "use_mask": "true",
          "prompt": "填充遮罩区域的头发，皮肤加一点磨皮，再加一点美颜"}
    if model_file:
        fd["flux_model"] = model_file
    if res_val != "":
        fd["redraw_max_side"] = res_val
    t0 = time.perf_counter()
    peak = gpu_used()
    err = None
    grown_b64 = None
    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=fd, timeout=300)
        elapsed = time.perf_counter() - t0
        peak = max(peak, gpu_used())
        j = r.json()
        if j.get("code") != 0:
            err = f"code={j.get('code')} {j.get('message', '')[:60]}"
        else:
            res = j.get("data", {}).get("results", [])
            if res and res[0].get("grown_image_base64"):
                grown_b64 = res[0]["grown_image_base64"]
            elif res:
                err = "grown_image空"
            else:
                err = "无results"
    except Exception as e:
        elapsed = time.perf_counter() - t0
        err = str(e)[:150]
    return {"elapsed": elapsed, "gpu_peak": peak, "grown_b64": grown_b64, "error": err}


def main():
    # 结果字典: results[(img, hair_num, column_title)] = {grown_path, elapsed, gpu_peak, error}
    results = {}
    total = len(IMGS) * len(HAIRSTYLES) * len(COLUMN_TITLES)
    idx = 0
    for mfile, res_list in MODEL_GROUPS:
        mname = os.path.basename(mfile)
        print(f"\n===== 切换到模型: {mname} =====", flush=True)
        # 等模型切换稳定
        time.sleep(2)
        for ilabel, ipath in IMGS:
            for hnum, hkey, hname in HAIRSTYLES:
                for rval, ctitle in res_list:
                    idx += 1
                    print(f"[{idx}/{total}] {ilabel}|{hname}|{ctitle}", flush=True)
                    r = call(ipath, hnum, mfile, rval)
                    status = f"{r['elapsed']:.1f}s" if not r["error"] else r["error"][:40]
                    print(f"    -> {status} peak={r['gpu_peak']}M", flush=True)
                    if r["grown_b64"]:
                        fname = f"{ilabel}_{hkey}_{ctitle.replace('@','_').replace('-','')}.jpg"
                        with open(OUT / fname, "wb") as gf:
                            gf.write(base64.b64decode(r["grown_b64"]))
                        r["grown_path"] = str(OUT / fname)
                    results[(ilabel, hnum, ctitle)] = r

    # 重组为15行
    rows = []
    for ilabel, ipath in IMGS:
        for hnum, hkey, hname in HAIRSTYLES:
            cells = []
            for ct in COLUMN_TITLES:
                r = results.get((ilabel, hnum, ct), {"error": "未跑"})
                cells.append({"title": ct, **{k: v for k, v in r.items() if k != "grown_b64"}})
            rows.append({"img": ilabel, "img_path": ipath,
                         "hair_num": hnum, "hair_key": hkey, "hair_name": hname,
                         "cells": cells})
    with open(OUT / "results.json", "w", encoding="utf-8") as f:
        json.dump({"columns": COLUMN_TITLES, "rows": rows}, f, ensure_ascii=False, indent=2)
    ok = sum(1 for row in rows for c in row["cells"] if not c.get("error"))
    print(f"\n✓ 完成: {ok}/{total} 成功 -> {OUT/'results.json'}", flush=True)


if __name__ == "__main__":
    main()