#!/usr/bin/env python3
"""
批量用 H3 ref2va (参考图模式) 生成 8 个 validation case。
- 用 H3 全参考模式 prompt 协议改写原始 prompt
- ref2va API + 标准 15 步
"""
import os, sys, time, json, subprocess, requests

KAP = "http://localhost:10588"
COMFY = "http://localhost:8188"
BASE = "/home/kai/文档/LTX/validition_v1"
OUT = "/home/kai/shared/validation_ref2va"
os.makedirs(OUT, exist_ok=True)

# Add enhancer path
sys.path.insert(0, "/data/workspace/kais-hermes-skills/skills/movie-experts/video-engine-adapter/scripts")
from h3_prompt_enhancer import build_local_template, get_system_prompt

CASES = [
    ("01", ["1.jpg", "2.jpg", "bg.png"], False),
    ("02", ["1.jpg", "2.jpg", "bg.png"], False),
    ("03", ["01.png", "02.png"], True),
    ("04", ["01.png", "02.png", "03.png"], False),
    ("05", ["1.png", "2.png", "5.png", "bg.png"], False),
    ("06", ["1.png", "2.png", "3.png", "bg.png"], True),
    ("07", ["1.jpg", "2.jpg", "bg.png"], False),
    ("08", ["1.jpg", "2.jpg", "bg.png"], False),
]

def read_prompt(case_dir):
    with open(os.path.join(case_dir, "prompt.txt")) as f:
        return f.read().strip()

def build_h3_prompt(raw_prompt, ref_files):
    """Build H3 ref2va prompt using local template + manual enhancement."""
    # Use the local template as a starting point
    template = build_local_template("全参考模式(Reference to Video)", raw_prompt, duration=8)
    
    # Build <Picture N> tags
    pic_tags = " ".join([f"<Picture {i+1}>" for i in range(len(ref_files))])
    
    # The template already has the right structure but missing Picture tags
    # We use the raw prompt directly as detailed_description with Picture tags
    # This is a simplified version - for production, use LLM API enhancement
    
    h3_prompt = f"""{pic_tags} {raw_prompt[:600]}

subject_definitions:
<Picture 1> 人物参考图，提供角色外观、服装、配饰和颜色特征。
"""

    if len(ref_files) >= 2:
        h3_prompt += f"<Picture 2> 人物参考图，提供第二个角色的外观、服装、配饰和颜色特征。\n"
    if len(ref_files) >= 3:
        h3_prompt += f"<Picture 3> 场景参考图，提供环境、光照和氛围。\n"
    if len(ref_files) >= 4:
        for i in range(3, len(ref_files)):
            h3_prompt += f"<Picture {i+1}> 补充参考图，提供额外角色或物体外观。\n"

    h3_prompt += f"""
summary:
[reference generation] 基于{len(ref_files)}张参考图生成8秒视频。

retention_analysis:
<Picture 1>（角色1外观）: fully_preserved - 保持参考图中人物的面部、发型、服装、配饰和颜色。
"""
    if len(ref_files) >= 2:
        h3_prompt += "<Picture 2>（角色2外观）: fully_preserved - 保持参考图中人物的面部、发型、服装、配饰和颜色。\n"
    if len(ref_files) >= 3:
        h3_prompt += "<Picture 3>（场景环境）: fully_preserved - 保持参考图中的场景布局、光照和氛围。\n"

    h3_prompt += f"""
detailed_description:
{raw_prompt[:600]}

overall_soundscape:
环境自然声音，包括脚步声、衣物摩擦声和自然风声。

non_diegetic_music:
N/A
"""
    return h3_prompt

def submit_ref2va(ref_paths, prompt, portrait):
    width, height = (768, 1344) if portrait else (1344, 768)
    data = {
        "prompt": prompt,
        "projectId": "999",
        "width": str(width),
        "height": str(height),
        "steps": "4",
        "turbo": "true",
        "refImageSize": "match",
    }
    files = []
    open_files = []
    for rp in ref_paths:
        fh = open(rp, "rb")
        open_files.append(fh)
        files.append(("refImages", (os.path.basename(rp), fh, "image/jpeg")))
    
    resp = requests.post(f"{KAP}/api/production/minimax-h3/ref2va", data=data, files=files, timeout=60)
    for fh in open_files:
        fh.close()
    return resp.json()

def poll_and_download(prompt_id, case_num, timeout=900):
    start = time.time()
    while time.time() - start < timeout:
        try:
            resp = requests.get(f"{KAP}/api/production/minimax-h3/status/{prompt_id}", timeout=15)
            d = resp.json()
            data = d.get("data", d)
            status = data.get("status", "")
            if status in ("success", "completed"):
                outputs = data.get("outputs", {})
                filename = None
                for nid, out in outputs.items():
                    if isinstance(out, dict):
                        for k in ("images", "videos", "gifs"):
                            if k in out and out[k]:
                                item = out[k][0]
                                filename = item.get("filename", "") if isinstance(item, dict) else str(item)
                                if filename.endswith((".mp4", ".webm")):
                                    break
                        if filename:
                            break
                if filename:
                    local = os.path.join(OUT, f"case{case_num}_ref2va.mp4")
                    url = f"{COMFY}/view?filename={filename}&type=output"
                    subprocess.run(["curl", "-sL", "-o", local, url], check=True, timeout=120)
                    sz = os.path.getsize(local) // 1024
                    elapsed = time.time() - start
                    print(f"✅ Case {case_num} DONE ({elapsed:.0f}s): {filename} → {sz}KB", flush=True)
                    return True
                else:
                    print(f"⚠️ Case {case_num} no filename: {json.dumps(data)[:300]}", flush=True)
                    return False
            elif status == "error":
                print(f"❌ Case {case_num} ERROR: {json.dumps(data)[:300]}", flush=True)
                return False
        except Exception as e:
            print(f"  poll err: {e}", flush=True)
        elapsed = int(time.time() - start)
        if elapsed % 60 == 0 and elapsed > 0:
            print(f"  ⏳ Case {case_num} {elapsed}s... ({status})", flush=True)
        time.sleep(5)
    print(f"❌ Case {case_num} TIMEOUT", flush=True)
    return False

for case_num, ref_files, portrait in CASES:
    case_dir = os.path.join(BASE, case_num)
    raw_prompt = read_prompt(case_dir)
    ref_paths = [os.path.join(case_dir, f) for f in ref_files]
    h3_prompt = build_h3_prompt(raw_prompt, ref_files)
    
    print(f"\n[{time.strftime('%H:%M:%S')}] Case {case_num}: ref2va ({len(ref_files)} refs, {'portrait' if portrait else 'landscape'})...", flush=True)
    print(f"  prompt length: {len(h3_prompt)} chars", flush=True)
    
    try:
        resp = submit_ref2va(ref_paths, h3_prompt, portrait)
        data = resp.get("data", resp)
        prompt_id = data.get("promptId", data.get("prompt_id", ""))
        if not prompt_id:
            print(f"  ⚠️ No promptId: {json.dumps(resp)[:400]}", flush=True)
            continue
        print(f"  promptId: {prompt_id}", flush=True)
        poll_and_download(prompt_id, case_num)
    except Exception as e:
        print(f"  ❌ Case {case_num}: {e}", flush=True)

print("\nALL_DONE", flush=True)