#!/usr/bin/env python3
"""
H3 Hybrid Validation: Turbo + i2va + refImages
标准 H3 prompt 协议转换 + 批量提交 8 个 case。

每个 case:
- 第一张图 = 首帧 (image field)
- 其余图 = refImages (参考图)
- H3 协议 prompt: <Picture 1> (首帧) + <Picture 2/3/...> (参考图)
- Turbo 4步 + T8 Dual-Clock 采样器
"""
import requests, json, os, time, sys

KAP = "http://localhost:10588"
OUTPUT_DIR = "/home/kai/shared/h3_hybrid_results"
os.makedirs(OUTPUT_DIR, exist_ok=True)

BASE_DIR = "/home/kai/文档/LTX/validition_v1"


def build_h3_prompt(raw_prompt: str, num_images: int) -> str:
    """
    按照H3提示词协议标准化：
    - <Picture 1> = 首帧 (i2va 首帧锚定 0.00s)
    - <Picture 2> ... <Picture N> = 参考图
    
    H3 协议要求：
    - I2VA 首行对齐首帧图
    - 参考标签引用图片
    - 保留原始场景描述
    """
    # 如果 prompt 已经有 <Picture> 标签，直接返回
    if "<Picture" in raw_prompt:
        return raw_prompt
    
    # 构建参考标签
    tags = " ".join([f"<Picture {i+1}>" for i in range(num_images)])
    
    # 清理原始 prompt 中的中文参考图描述（"参考图1：..." 这类）
    lines = raw_prompt.strip().split("\n")
    clean_lines = []
    for line in lines:
        # 跳过"参考图N："开头的中文描述行
        stripped = line.strip()
        if stripped.startswith("参考图") or stripped.startswith("人物形象") or stripped.startswith("提示词："):
            continue
        clean_lines.append(line)
    
    clean_prompt = "\n".join(clean_lines).strip()
    
    # H3 协议格式：首帧对齐 + 参考标签 + 场景描述
    h3_prompt = f"{tags}\n{clean_prompt}"
    
    return h3_prompt


def submit_case(case_id: str, case_dir: str) -> dict:
    """提交一个 case 到 KAP generate API (hybrid mode)."""
    print(f"\n{'='*60}")
    print(f"Case {case_id}:")
    
    # 读 prompt
    prompt_file = os.path.join(case_dir, "prompt.txt")
    if not os.path.exists(prompt_file):
        prompt_file = os.path.join(case_dir, "prompt_nsfw.txt")
    if not os.path.exists(prompt_file):
        return {"case": case_id, "status": "skip", "error": "no prompt"}
    
    with open(prompt_file) as f:
        raw_prompt = f.read().strip()
    
    # 找图片
    imgs = sorted([f for f in os.listdir(case_dir) if f.endswith((".png", ".jpg", ".webp"))])
    if not imgs:
        return {"case": case_id, "status": "skip", "error": "no images"}
    
    # 第一张 = 首帧，其余 = refImages
    first_frame = os.path.join(case_dir, imgs[0])
    ref_images = [os.path.join(case_dir, imgs[i]) for i in range(1, min(len(imgs), 4))]  # max 3 refs
    
    # H3 prompt 标准化
    total_imgs = 1 + len(ref_images)  # 首帧 + 参考图
    h3_prompt = build_h3_prompt(raw_prompt, total_imgs)
    
    print(f"  Images: first_frame={imgs[0]}, refs={[os.path.basename(r) for r in ref_images]}")
    print(f"  Prompt: {len(h3_prompt)} chars (raw={len(raw_prompt)})")
    print(f"  H3 tags: {' '.join([f'<Picture {i+1}>' for i in range(total_imgs)])}")
    
    # 提交到 KAP generate API
    url = f"{KAP}/api/production/minimax-h3/generate"
    data = {
        "projectId": "0",
        "mode": "i2va",
        "prompt": h3_prompt,
        "width": "960",
        "height": "512",
        "length": "124",  # ~5s, n%17==5 → 124
        "turbo": "true",
        "profile": "turbo",
        "filenamePrefix": f"h3_hybrid_case{case_id}",
    }
    
    files = [("image", (imgs[0], open(first_frame, "rb")))]
    for r in ref_images:
        files.append(("refImages", (os.path.basename(r), open(r, "rb"))))
    
    print(f"  Submitting to KAP (i2va + {len(ref_images)} refImages, turbo)...")
    try:
        resp = requests.post(url, data=data, files=files, timeout=600)
    except Exception as e:
        return {"case": case_id, "status": "error", "error": f"request failed: {e}"}
    finally:
        for _, f in files:
            f[1].close()
    
    if resp.status_code != 200:
        err = resp.text[:300]
        print(f"  ❌ HTTP {resp.status_code}: {err}")
        return {"case": case_id, "status": "error", "error": f"HTTP {resp.status_code}: {err}"}
    
    r = resp.json()
    if r.get("code") != 200:
        err = json.dumps(r, ensure_ascii=False)[:300]
        print(f"  ❌ API error: {err}")
        return {"case": case_id, "status": "error", "error": err}
    
    result = r.get("data", {})
    video_url = result.get("videoUrl", "")
    print(f"  ✅ videoUrl: {video_url}")
    
    # 下载视频
    if video_url:
        # 从 KAP 获取文件
        dl_url = f"http://localhost:10588{video_url}" if video_url.startswith("/") else video_url
        try:
            dl_resp = requests.get(dl_url, timeout=60)
            if dl_resp.status_code == 200:
                outpath = os.path.join(OUTPUT_DIR, f"case{case_id}_hybrid.mp4")
                with open(outpath, "wb") as f:
                    f.write(dl_resp.content)
                print(f"  📁 Saved: {outpath} ({len(dl_resp.content)//1024}KB)")
                return {"case": case_id, "status": "ok", "file": outpath, "videoUrl": video_url}
        except Exception as e:
            print(f"  ⚠️ Download failed: {e}")
    
    return {"case": case_id, "status": "ok", "videoUrl": video_url}


def main():
    cases = sorted([d for d in os.listdir(BASE_DIR) if os.path.isdir(os.path.join(BASE_DIR, d))])
    
    # 支持只跑指定 case
    if len(sys.argv) > 1:
        cases = [c for c in cases if c in sys.argv[1:]]
    
    print(f"Running {len(cases)} cases: {cases}")
    
    results = []
    for case_id in cases:
        case_dir = os.path.join(BASE_DIR, case_id)
        result = submit_case(case_id, case_dir)
        results.append(result)
    
    # Summary
    print(f"\n{'='*60}")
    print("SUMMARY:")
    for r in results:
        status = r["status"]
        case = r["case"]
        if status == "ok":
            print(f"  Case {case}: ✅ {r.get('videoUrl', r.get('file', ''))}")
        else:
            print(f"  Case {case}: ❌ {r.get('error', status)}")
    
    ok = sum(1 for r in results if r["status"] == "ok")
    print(f"\n{ok}/{len(results)} cases succeeded")
    
    # Save results JSON
    with open(os.path.join(OUTPUT_DIR, "results.json"), "w") as f:
        json.dump(results, f, indent=2, ensure_ascii=False)
    
    print("ALL_DONE")


if __name__ == "__main__":
    main()