#!/usr/bin/env python3
"""
直接构建 ComfyUI API workflow 提交 ref2va (绕过 KAP ref2va 路由的 autogrow bug)。
所有 8 个 validation case，Turbo 4 步。
"""
import os, sys, time, json, subprocess, urllib.request

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

# H3 model constants
MODEL = "minimax_h3_ref2va_pruned_int8_convrot.safetensors"
TURBO_LORA = "minimax_h3_turbo_4step_ema_original_comfyui.safetensors"
CLIP = "qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors"
VIDEO_VAE = "minimax_h3_video_vae_fp16.safetensors"
AUDIO_VAE = "minimax_h3_audio_vae_fp32.safetensors"

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

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

def build_workflow(case_num, ref_filenames, prompt, width, height, seed=42):
    """Build ComfyUI API workflow for ref2va with T8 node."""
    import random
    if seed is None:
        seed = random.randint(1, 2**31)
    
    # Copy ref images to container
    case_dir = os.path.join(BASE, case_num)
    container_filenames = []
    for rf in ref_filenames:
        local = os.path.join(case_dir, rf)
        ext = os.path.splitext(rf)[1]
        cname = f"val_{case_num}_{rf.replace('/', '_')}"
        subprocess.run(["docker", "cp", local, f"comfyui-primary:/root/ComfyUI/input/{cname}"],
                       check=True, timeout=30)
        container_filenames.append(cname)
    
    num_refs = len(container_filenames)
    pic_tags = " ".join([f"<Picture {i+1}>" for i in range(num_refs)])
    full_prompt = f"{pic_tags} {prompt[:600]}"
    
    # Length: 5s @ 24fps → align to n%17==5 → 125 frames
    length = 125
    
    # Build LoadImage nodes for refs: node 14, 141, 142...
    load_image_nodes = {}
    ref_image_refs = []
    for i, fname in enumerate(container_filenames):
        nid = 14 + i if i == 0 else 140 + i  # 14, 141, 142, 143
        load_image_nodes[str(nid)] = {
            "class_type": "LoadImage",
            "inputs": {"image": fname}
        }
        ref_image_refs.append([str(nid), 0])
    
    wf = {
        # Loaders
        "10": {"class_type": "CLIPLoader", "inputs": {"clip_name": CLIP, "type": "minimax"}},
        "11": {"class_type": "VAELoader", "inputs": {"vae_name": VIDEO_VAE}},
        "12": {"class_type": "UNETLoader", "inputs": {"unet_name": MODEL, "weight_dtype": "default"}},
        "13": {"class_type": "VAELoader", "inputs": {"vae_name": AUDIO_VAE}},
        
        # Turbo LoRA
        "14_lora": {
            "class_type": "LoraLoaderBypassModelOnly",
            "inputs": {
                "model": ["12", 0],
                "lora_name": TURBO_LORA,
                "strength_model": 1.0,
            }
        },
    }
    
    # Add LoadImage nodes
    wf.update(load_image_nodes)
    
    # Node 20: T8 Conditioning - use ref_images as list of [nodeId, 0]
    wf["20"] = {
        "class_type": "MiniMaxH3AudioConditioningT8",
        "inputs": {
            "clip": ["10", 0],
            "video_vae": ["11", 0],
            "audio_vae": ["13", 0],
            "prompt": full_prompt,
            "width": width,
            "height": height,
            "length": length,
            "task_type": "auto",
            "audio_mode": "native",
            "audio_denoise_strength": 0.35,
            "add_source_as_reference": False,
            "prompt_primary_audio_ordinal": 1,
            "strict_prompt_tags": False,  # Disable strict to avoid validation issues
            "ref_image_size": "match",
            "reference_video_policy": "official_2_to_15s",
            "ref_images": ref_image_refs,
        }
    }
    
    # Sampler
    wf["30"] = {
        "class_type": "MiniMaxH3DualClockSamplerT8",
        "inputs": {
            "model": ["14_lora", 0],
            "av_latent": ["20", 1],
            "steps": 4,
            "shift_video": 12,
            "shift_audio": 3,
        }
    }
    wf["31"] = {"class_type": "RandomNoise", "inputs": {"noise_seed": seed}}
    wf["32"] = {"class_type": "BasicGuider", "inputs": {"model": ["30", 0], "conditioning": ["20", 0]}}
    wf["33"] = {
        "class_type": "SamplerCustomAdvanced",
        "inputs": {"noise": ["31", 0], "guider": ["32", 0], "sampler": ["30", 1], "sigmas": ["30", 2], "latent_image": ["20", 1]}
    }
    wf["40"] = {
        "class_type": "MiniMaxH3AVDecodeT8",
        "inputs": {"av_latent": ["33", 0], "video_vae": ["11", 0], "audio_vae": ["13", 0]}
    }
    wf["42"] = {
        "class_type": "CreateVideo",
        "inputs": {"images": ["40", 0], "fps": 24, "audio": ["40", 1]}
    }
    wf["50"] = {
        "class_type": "SaveVideo",
        "inputs": {"video": ["42", 0], "filename_prefix": f"val_ref2va_case{case_num}", "format": "mp4", "codec": "auto"}
    }
    
    return wf

def submit(wf):
    req = urllib.request.Request(
        f"{COMFY}/prompt",
        data=json.dumps({"prompt": wf}).encode(),
        headers={'Content-Type': 'application/json'})
    with urllib.request.urlopen(req, timeout=30) as r:
        return json.load(r)

def wait_and_download(pid, case_num, timeout=600):
    start = time.time()
    while time.time() - start < timeout:
        try:
            with urllib.request.urlopen(f"{COMFY}/history/{pid}", timeout=15) as r:
                d = json.load(r)
            if pid in d:
                st = d[pid].get('status', {})
                if st.get('completed'):
                    files = []
                    for nid, out in d[pid].get('outputs', {}).items():
                        for k in ('gifs', 'videos', 'images'):
                            if k in out:
                                for item in out[k]:
                                    fn = item.get('filename', '')
                                    if fn.endswith(('.mp4', '.webm')):
                                        files.append(fn)
                    if files:
                        elapsed = time.time() - start
                        local = os.path.join(OUT, f"case{case_num}_ref2va.mp4")
                        url = f"{COMFY}/view?filename={files[0]}&type=output"
                        subprocess.run(["curl", "-sL", "-o", local, url], check=True, timeout=120)
                        sz = os.path.getsize(local) // 1024
                        print(f"✅ Case {case_num} DONE ({elapsed:.0f}s): {files[0]} → {sz}KB", flush=True)
                        return True
                for m in st.get('messages', []):
                    if m[0] == 'execution_error':
                        msg = m[1].get('exception_message', '')[:300]
                        print(f"❌ Case {case_num} ERROR: {msg}", 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...", flush=True)
        time.sleep(5)
    print(f"❌ Case {case_num} TIMEOUT", flush=True)
    return False

for case_num, ref_files, width, height in CASES:
    case_dir = os.path.join(BASE, case_num)
    raw_prompt = read_prompt(case_dir)
    
    print(f"\n[{time.strftime('%H:%M:%S')}] Case {case_num}: building workflow ({len(ref_files)} refs, {width}x{height})...", flush=True)
    
    wf = build_workflow(case_num, ref_files, raw_prompt, width, height)
    
    print(f"  Submitting...", flush=True)
    try:
        resp = submit(wf)
        if 'error' in resp and resp['error']:
            print(f"  ❌ Submit failed: {json.dumps(resp['error'])[:400]}", flush=True)
            continue
        pid = resp.get('prompt_id', '')
        print(f"  promptId: {pid}", flush=True)
        wait_and_download(pid, case_num)
    except Exception as e:
        print(f"  ❌ Case {case_num}: {e}", flush=True)

print("\nALL_DONE", flush=True)