#!/usr/bin/env python3
"""9B vs 4B 模型对比：3 张图 × 5 种发型 = 15 张对比图。

第一次运行：PLAN_TAG=A  （9B 模型，需先切工作流到 .bak）
第二次运行：PLAN_TAG=Bplus（4B 模型，需切回 4B 工作流）
"""
import json, os, time, base64, subprocess, urllib.request
from datetime import datetime

API_BASE = "http://127.0.0.1:8187"
TOKEN = "dev-shared-secret-2026"
IMAGES = [
    "/home/ubuntu/hair/image/girl_img/girl1.jpg",
    "/home/ubuntu/hair/image/girl_img/girl7.jpg",
    "/home/ubuntu/hair/image/girl_img/girl13.jpg",
]
PLAN_TAG = os.getenv("PLAN_TAG", "A")
OUT_DIR = f"/home/ubuntu/hair/benchmark_out/9b_vs_4b_{PLAN_TAG}"
RESULT_JSON = os.path.join(OUT_DIR, "results.json")
REPORT_HTML = os.path.join(OUT_DIR, "report.html")
HAIR_STYLES = [(1, "椭圆"), (2, "花瓣"), (3, "心形"), (4, "直线"), (5, "波浪")]
WARMUP = os.getenv("WARMUP", "1") == "1"  # 第一次调用做预热

os.makedirs(OUT_DIR, exist_ok=True)


def get_vram():
    try:
        out = subprocess.check_output(
            ["nvidia-smi", "--query-gpu=memory.used,memory.free",
             "--format=csv,noheader,nounits"], text=True, timeout=5).strip()
        used, free = out.split(",")
        return int(used.strip()), int(free.strip())
    except Exception:
        return -1, -1


def _multipart(fields, files=None):
    boundary = "----CmpTest" + str(int(time.time() * 1000))
    parts = []
    for k, v in fields.items():
        parts.append(f"--{boundary}\r\n".encode())
        parts.append(f'Content-Disposition: form-data; name="{k}"\r\n\r\n'.encode())
        parts.append(str(v).encode())
        parts.append(b"\r\n")
    if files:
        for field_name, (filename, data, mime) in files.items():
            parts.append(f"--{boundary}\r\n".encode())
            parts.append(f'Content-Disposition: form-data; name="{field_name}"; filename="{filename}"\r\n'.encode())
            parts.append(f"Content-Type: {mime}\r\n\r\n".encode())
            parts.append(data)
            parts.append(b"\r\n")
    parts.append(f"--{boundary}--\r\n".encode())
    return b"".join(parts), boundary


def call_iface2(image_path, hair_style):
    with open(image_path, "rb") as f:
        img_data = f.read()
    fields = {"gender": "female", "hair_style": str(hair_style),
              "prompt": "填充遮罩区域的头发"}
    body, boundary = _multipart(
        fields, {"image_file": (os.path.basename(image_path), img_data, "image/jpeg")})
    req = urllib.request.Request(f"{API_BASE}/api/v1/hair/grow", data=body, method="POST")
    req.add_header("Content-Type", f"multipart/form-data; boundary={boundary}")
    req.add_header("X-Internal-Token", TOKEN)
    t0 = time.time()
    try:
        with urllib.request.urlopen(req, timeout=600) as resp:
            data = json.loads(resp.read())
            return data, time.time() - t0, None
    except Exception as e:
        return {}, time.time() - t0, str(e)


def extract_image(data, img_idx, hs):
    """从响应提取 grown_image_base64 并保存为 jpg。"""
    if not data or not data.get("data"):
        return None
    items = data["data"].get("results") or []
    if not items:
        return None
    b64 = items[0].get("grown_image_base64") or ""
    if b64.startswith("data:"):
        b64 = b64.split(",", 1)[1]
    if not b64:
        return None
    fname = f"img{img_idx}_style{hs}.jpg"
    path = os.path.join(OUT_DIR, fname)
    with open(path, "wb") as f:
        f.write(base64.b64decode(b64))
    return path


def run_test():
    results = []
    print(f"9B vs 4B 对比测试 — Plan={PLAN_TAG}")
    print(f"图片: {[os.path.basename(p) for p in IMAGES]}")
    print(f"发型: {[(s,n) for s,n in HAIR_STYLES]}")
    print(f"总调用: {len(IMAGES)*len(HAIR_STYLES)}\n")

    # 预热请求（避免第一次冷启动计入统计）
    if WARMUP:
        print("[warmup] 预热请求 (girl13, style=1) ...")
        t0 = time.time()
        _, warmup_time, _ = call_iface2(IMAGES[2], 1)
        print(f"  warmup: {warmup_time:.1f}s\n")

    for img_idx, img_path in enumerate(IMAGES, start=1):
        img_name = os.path.basename(img_path)
        for hs, hs_name in HAIR_STYLES:
            print(f"[img{img_idx}/{len(IMAGES)}] {img_name} style={hs}({hs_name}) ...")

            vram_before, _ = get_vram()
            data, elapsed, error = call_iface2(img_path, hs)
            vram_after, _ = get_vram()

            code = data.get("code", -1) if data else -1
            ok = (code == 0)
            img_saved = extract_image(data, img_idx, hs) if ok else None

            record = {
                "plan": PLAN_TAG,
                "image_idx": img_idx,
                "image_name": img_name,
                "image_path": img_path,
                "hair_style": hs,
                "hair_style_name": hs_name,
                "timestamp": datetime.now().strftime("%H:%M:%S"),
                "elapsed_s": round(elapsed, 2),
                "success": ok,
                "code": code,
                "error": error,
                "vram_before_mb": vram_before,
                "vram_after_mb": vram_after,
                "vram_delta_mb": vram_after - vram_before,
                "saved_image_path": img_saved,
            }
            results.append(record)

            status = "✓" if ok else "✗"
            print(f"  → {status} code={code} time={elapsed:.1f}s "
                  f"vram={vram_before}→{vram_after}MB (Δ{vram_after-vram_before:+d}) "
                  f"img={'saved' if img_saved else 'none'}")

            with open(RESULT_JSON, "w") as f:
                json.dump(results, f, indent=2, ensure_ascii=False)

    print(f"\n✅ 完成！{RESULT_JSON}")
    print(f"📊 成功: {sum(1 for r in results if r['success'])}/{len(results)}")
    if any(r['success'] for r in results):
        avg = sum(r['elapsed_s'] for r in results if r['success']) / sum(1 for r in results if r['success'])
        print(f"⏱  平均: {avg:.2f}s")


if __name__ == "__main__":
    run_test()