#!/usr/bin/env python3
"""继续生成 Case 04-08 (01-03 已完成)"""
import os, time, json, subprocess, requests

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

CASES = [
    ("04", "01.png", "16:9", 8),
    ("05", "1.png", "16:9", 8),
    ("06", "1.png", "9:16", 8),
    ("07", "1.jpg", "16:9", 8),
    ("08", "1.jpg", "16:9", 8),
]

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

def submit_i2va(ff_path, prompt, aspect, duration):
    data = {
        "prompt": prompt[:800],
        "projectId": "999",
        "turbo": "true",
        "aspectRatio": aspect,
        "duration": str(duration),
        "steps": "4",
    }
    with open(ff_path, "rb") as f:
        files = {"firstFrame": (os.path.basename(ff_path), f, "image/jpeg")}
        resp = requests.post(f"{KAP}/api/production/minimax-h3/i2va", data=data, files=files, timeout=60)
    return resp.json()

def poll_and_download(prompt_id, case_num, timeout=600):
    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"):
                # Extract filename from ComfyUI outputs format
                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:
                    # Download via ComfyUI /view
                    local = os.path.join(OUT, f"case{case_num}_h3turbo.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} completed but no video filename found", 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, ff_file, aspect, duration in CASES:
    case_dir = os.path.join(BASE, case_num)
    ff_path = os.path.join(case_dir, ff_file)
    prompt = read_prompt(case_dir)
    
    print(f"\n[{time.strftime('%H:%M:%S')}] Case {case_num}: submitting...", flush=True)
    try:
        resp = submit_i2va(ff_path, prompt, aspect, duration)
        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)