#!/usr/bin/env python3
"""接口2 female 批量对比测试：4图 × 5发型 × 3分辨率档位 = 60次请求。
串行执行，记录耗时与成败，结果图按 图_发型_档位 命名保存。
用法: python3 batch_test.py
支持断点续跑（progress.json）；新增档位时只会补跑未完成项。
"""
import base64, csv, json, os, sys, time, traceback
import requests

API = "http://127.0.0.1:8187/api/v1/hair/grow"
TOKEN = "dev-shared-secret-2026"
TIMEOUT = 600
OUT = "/home/ubuntu/hair/image/compare_test/out"
PROGRESS = "/home/ubuntu/hair/image/compare_test/progress.json"

IMG_DIR = "/home/ubuntu/hair/image"
IMAGES = [
    ("girl2",  f"{IMG_DIR}/girl_img/girl2.jpg"),
    ("girl5",  f"{IMG_DIR}/girl_img/girl5.jpg"),
    ("qwer",   f"{IMG_DIR}/qwer.jpg"),
    ("asdf",   f"{IMG_DIR}/asdf.jpg"),
]
# female: 1=ellipse 2=flower 3=heart 4=straight 5=wave
STYLES = [
    (1, "ellipse"), (2, "flower"), (3, "heart"), (4, "straight"), (5, "wave"),
]
# 三档: 默认896 / 显式1024 / 原图直送0
SIDES = [
    ("default896", None),   # 不传 → 后端默认896
    ("side1024",   1024),   # 长边压到 1024
    ("origin0",    0),      # 原图直送
]
SIDE_LABEL = {
    "default896": "默认896",
    "side1024": "1024",
    "origin0": "原图0",
}

def load_progress():
    if os.path.exists(PROGRESS):
        try:
            return json.load(open(PROGRESS))
        except Exception:
            pass
    return {"done": [], "results": []}

def save_progress(prog):
    json.dump(prog, open(PROGRESS, "w"), ensure_ascii=False, indent=1)

def run_one(img_name, img_path, style_idx, style_key, side_name, side_val):
    """跑单次请求，返回 dict 结果。"""
    key = f"{img_name}_{style_key}_{side_name}"
    with open(img_path, "rb") as f:
        img_b64 = base64.b64encode(f.read()).decode()
    data = {
        "image_base64": "data:image/jpeg;base64," + img_b64,
        "gender": "female",
        "hair_style": str(style_idx),
    }
    if side_val is not None:
        data["redraw_max_side"] = str(side_val)
    t0 = time.time()
    rec = {"key": key, "img": img_name, "style": style_key, "style_idx": style_idx,
           "side": side_name, "side_val": side_val, "ok": False,
           "elapsed": 0.0, "err": "", "out_w": 0, "out_h": 0}
    try:
        r = requests.post(API, data=data, headers={"X-Internal-Token": TOKEN}, timeout=TIMEOUT)
        rec["elapsed"] = round(time.time() - t0, 1)
        d = r.json()
        if d.get("code") != 0:
            rec["err"] = f"code={d.get('code')} {d.get('message','')}"[:200]
            return rec
        results = (d.get("data") or {}).get("results") or []
        if not results:
            rec["err"] = "空结果"
            return rec
        it = results[0]
        grown = it.get("grown_image_base64")
        if not grown:
            rec["err"] = "无生发图(grown_png=None, 重绘失败)"
            return rec
        raw = base64.b64decode(grown)
        from PIL import Image
        import io as _io
        im = Image.open(_io.BytesIO(raw))
        rec["out_w"], rec["out_h"] = im.size
        out_path = f"{OUT}/{key}.jpg"
        with open(out_path, "wb") as fo:
            fo.write(raw)
        rec["ok"] = True
        rec["bytes"] = len(raw)
    except requests.exceptions.Timeout:
        rec["elapsed"] = round(time.time() - t0, 1)
        rec["err"] = f"超时(>{TIMEOUT}s)"
    except Exception as e:
        rec["elapsed"] = round(time.time() - t0, 1)
        rec["err"] = f"{type(e).__name__}: {str(e)[:180]}"
    return rec

def main():
    os.makedirs(OUT, exist_ok=True)
    prog = load_progress()
    done_keys = set(prog["done"])
    total = len(IMAGES) * len(STYLES) * len(SIDES)
    print(f"=== 接口2 female 批量对比测试 ===")
    print(f"矩阵: {len(IMAGES)}图 × {len(STYLES)}发型 × {len(SIDES)}档 = {total} 次")
    print(f"已跳过 {len(done_keys)} 个已完成项\n")

    idx = 0
    for img_name, img_path in IMAGES:
        for style_idx, style_key in STYLES:
            for side_name, side_val in SIDES:
                idx += 1
                key = f"{img_name}_{style_key}_{side_name}"
                if key in done_keys:
                    print(f"[{idx}/{total}] ⏭  跳过已完成 {key}")
                    continue
                print(f"[{idx}/{total}] ▶ {key} (side={side_val})...", end=" ", flush=True)
                rec = run_one(img_name, img_path, style_idx, style_key, side_name, side_val)
                prog["results"].append(rec)
                prog["done"].append(key)
                save_progress(prog)
                if rec["ok"]:
                    print(f"✅ {rec['elapsed']}s  出图 {rec['out_w']}x{rec['out_h']} ({rec['bytes']}B)")
                else:
                    print(f"❌ {rec['elapsed']}s  {rec['err']}")
                time.sleep(2)  # 串行间隔，避免队列堆积

    # 汇总
    print("\n" + "="*70)
    print("汇总报告")
    print("="*70)
    write_report(prog["results"])
    print(f"\n结果图: {OUT}/")
    print(f"CSV报告: {OUT}/report.csv")
    print(f"JSON明细: {OUT}/report.json")

def write_report(results):
    # CSV
    csv_path = f"{OUT}/report.csv"
    with open(csv_path, "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["img","style","side","ok","elapsed_s","out_w","out_h","err"])
        for r in results:
            w.writerow([r["img"],r["style"],r["side"],r["ok"],r["elapsed"],
                        r["out_w"],r["out_h"],r["err"]])
    # JSON
    json.dump(results, open(f"{OUT}/report.json","w"), ensure_ascii=False, indent=1)
    # 控制台速度对比表
    print("\n--- 速度对比（秒，✓=成功 ✗=失败）---")
    headers = ["图", "发型"] + [SIDE_LABEL[s] for s, _ in SIDES]
    print(f"{headers[0]:<8}{headers[1]:<10}" + "".join(f"{h:>12}" for h in headers[2:]))
    for img_name, _ in IMAGES:
        for _, style_key in STYLES:
            cells = []
            for side_name, _ in SIDES:
                r = next((x for x in results if x["img"]==img_name and x["style"]==style_key and x["side"]==side_name), None)
                cells.append(f"{r['elapsed']}{'✓' if r['ok'] else '✗'}" if r else "-")
            print(f"{img_name:<8}{style_key:<10}" + "".join(f"{c:>12}" for c in cells))

if __name__ == "__main__":
    main()