#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""接口2女性 花瓣形 单发型 4模型×3分辨率×3图×3次 矩阵测试。

调用本机 hair-worker (:8187) 的 /api/v1/hair/grow，gender=female, hair_style=2(花瓣形)。
每次记录：生发图、耗时、显存峰值。结果图存到 benchmark_out/matrix/，最后生成 HTML 报告。
"""
import base64
import json
import os
import subprocess
import sys
import time
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/matrix")
OUT.mkdir(parents=True, exist_ok=True)

# 4 模型 × 3 分辨率 × 3 图 × 3 次
MODELS = [
    ("4b-fp8", "flux-2-klein-4b-fp8.safetensors"),
    ("9b-fp8", "flux2.0/flux-2-klein-9b-fp8.safetensors"),
    ("9b-Q5", "flux-2-klein-9b-Q5_K_M.gguf"),
    ("9b-Q4", "flux-2-klein-9b-Q4_K_M.gguf"),
]
RES = [("orig", "0"), ("640", "640"), ("896", "896")]
IMGS = [
    ("asdf", "/home/ubuntu/hair/image/asdf.jpg"),
    ("qwer", "/home/ubuntu/hair/image/qwer.jpg"),
    ("girl5", "/home/ubuntu/hair/image/girl_img/girl5.jpg"),
]
REPEAT = 3


def gpu_used():
    """返回当前显存已用 MiB。"""
    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, model_file, res_val):
    """调一次接口2。返回 dict: ok/elapsed/grown_path/gpu_peak/error。"""
    fd = {
        "gender": "female",
        "hair_style": "2",  # 花瓣形
        "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_path = 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','')}"
        else:
            res = j.get("data", {}).get("results", [])
            if res and res[0].get("grown_image_base64"):
                grown_path = OUT / f"tmp_grown.jpg"
                with open(grown_path, "wb") as gf:
                    gf.write(base64.b64decode(res[0]["grown_image_base64"]))
            elif res:
                err = "grown_image_base64 为空"
            else:
                err = "无 results"
    except Exception as e:
        elapsed = time.perf_counter() - t0
        err = str(e)[:200]
    return {"elapsed": elapsed, "gpu_peak": peak, "grown_path": str(grown_path) if grown_path else None, "error": err}


def main():
    results = []  # 每元素一个组合
    total = len(MODELS) * len(RES) * len(IMGS) * REPEAT
    idx = 0
    for mlabel, mfile in MODELS:
        for rlabel, rval in RES:
            for ilabel, ipath in IMGS:
                # 一个组合：3 次
                runs = []
                for rep in range(REPEAT):
                    idx += 1
                    print(f"[{idx}/{total}] {mlabel} | res={rlabel} | {ilabel} | rep{rep+1}", flush=True)
                    r = call(ipath, mfile, rval)
                    print(f"    -> {r['elapsed']:.1f}s  peak={r['gpu_peak']}MiB  err={r['error']}", flush=True)
                    # 存每次的生发图
                    if r["grown_path"]:
                        save_to = OUT / f"{mlabel}_{rlabel}_{ilabel}_r{rep+1}.jpg"
                        os.replace(r["grown_path"], save_to)
                        r["grown_path"] = str(save_to)
                    runs.append(r)
                results.append({
                    "model": mlabel, "model_file": mfile,
                    "res": rlabel, "res_val": rval,
                    "img": ilabel, "img_path": ipath,
                    "runs": runs,
                })
    # 存原始数据
    with open(OUT / "results.json", "w", encoding="utf-8") as f:
        json.dump(results, f, ensure_ascii=False, indent=2)
    print(f"\n✓ 全部完成，原始数据 -> {OUT/'results.json'}", flush=True)


if __name__ == "__main__":
    main()