#!/usr/bin/env python3
"""RefMods probe: three-arm A/B on Case01 (RN+WV double character, 1216x672, ~5s).
Arm N: native ref2va with 2 refImages (1.jpg RN + 2.jpg WV) + bg prompt only in text
Arm E: RefMod encode (case01_rn_encode + case01_wv_encode)
Arm T: RefMod training 16x16 (case01_rn_train16 + case01_wv_train16)
Same seed, same prompt, same sampler settings.
"""
import json, urllib.request, time, sys, subprocess

BASE = "http://127.0.0.1:8188"
opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))

SEED = 20260908
WIDTH, HEIGHT, LENGTH = 1216, 672, 124  # 5.17s @24fps, div32-ish ok (1216/32=38, 672/32=21)

PROMPT = """subject_definitions:
<Subject 1> RN, a handsome man with long dark-brown half-tied hair, black dragon hairpin with pearl tassel, wine-red wide-sleeve robe with golden phoenix embroidery, black fox-fur cloak, jade belt
<Subject 2> WV, a gentle woman with black updo hair, jade-white flower hairpiece, dark-green fishtail gown with golden floral embroidery, translucent wide-sleeve shawl, jade pendant necklace

integrated_multimodal_description:
[Shot 1] Cinematic live-action. Rainy Jiangnan water town at dusk, wet stone-paved lane with mirror reflections, white walls and black tiles, warm lanterns, a great willow tree, misty green mountains. <Subject 2> WV walks toward camera holding her skirt, expression turning from cool to a soft smile. <Subject 1> RN stands facing her, smiling gently. They stop face to face on the stone lane, WV nods slightly with a faint blush. Medium two-shot, shallow depth of field, 35mm film look."""

NEG = "worst quality, blurry, watermark, text, subtitles, extra fingers, deformed hands, distorted face"

def post(wf, timeout=60):
    req = urllib.request.Request(BASE + "/prompt", data=json.dumps({"prompt": wf, "client_id": "refmod-probe"}).encode(),
                                 headers={"Content-Type": "application/json"})
    r = json.loads(opener.open(req, timeout=timeout).read())
    if r.get("node_errors"):
        raise RuntimeError("node_errors: " + json.dumps(r["node_errors"])[:800])
    return r["prompt_id"]

def wait(prompt_id, timeout=2400):
    t0 = time.time()
    while time.time() - t0 < timeout:
        try:
            h = json.loads(opener.open(f"{BASE}/history/{prompt_id}", timeout=10).read())
        except Exception:
            h = {}
        if prompt_id in h:
            st = h[prompt_id].get("status", {})
            msgs = st.get("messages", [])
            errs = [m for m in msgs if m[0] == "execution_error"]
            return st.get("status_str"), errs, h[prompt_id].get("outputs", {})
        time.sleep(8)
    return "TIMEOUT", [], {}

def loadimg(name):
    return f'"__LOAD__{name}"'

def build_arm(arm):
    wf = {}
    # shared: clip, vaes, unet
    wf["10"] = {"class_type": "CLIPLoader", "inputs": {"clip_name": "qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", "type": "minimax"}}
    wf["11"] = {"class_type": "VAELoader", "inputs": {"vae_name": "minimax_h3_video_vae_fp16.safetensors"}}
    wf["12"] = {"class_type": "VAELoader", "inputs": {"vae_name": "minimax_h3_audio_vae_fp32.safetensors"}}
    wf["13"] = {"class_type": "UNETLoader", "inputs": {"unet_name": "minimax_h3_ref2va_pruned_int8_convrot.safetensors", "weight_dtype": "default"}}

    if arm == "N":
        # native: 2 refImages via ReferenceToVideo
        wf["20"] = {"class_type": "LoadImage", "inputs": {"image": "kais_validation/01/1.jpg"}}
        wf["21"] = {"class_type": "LoadImage", "inputs": {"image": "kais_validation/01/2.jpg"}}
        wf["30"] = {"class_type": "MiniMaxH3ReferenceToVideo", "inputs": {
            "clip": ["10", 0], "vae": ["11", 0], "audio_vae": ["12", 0],
            "prompt": PROMPT, "width": WIDTH, "height": HEIGHT, "length": LENGTH,
            "ref_image_size": "match",
            "ref_images.ref_image_1": ["20", 0],
            "ref_images.ref_image_2": ["21", 0],
        }}
        pos = ["30", 0]
        latent_src = ["30", 1]
    else:
        suffix = "encode" if arm == "E" else "train16"
        # RefMod arms: conditioning from ImageToVideo-like path? No — RefModApply needs CONDITIONING.
        # Use ReferenceToVideo WITHOUT ref images (text-only cond+latent), then Apply injects mods.
        wf["30"] = {"class_type": "MiniMaxH3ReferenceToVideo", "inputs": {
            "clip": ["10", 0], "vae": ["11", 0], "audio_vae": ["12", 0],
            "prompt": PROMPT, "width": WIDTH, "height": HEIGHT, "length": LENGTH,
            "ref_image_size": "match",
        }}
        loader_inputs = {"load": True, "show_info": False}
        for slot in range(1, 9):
            loader_inputs[f"mod_{slot}"] = "(none)"
            loader_inputs[f"strength_{slot}"] = 1.0
            loader_inputs[f"copies_{slot}"] = 1
        loader_inputs["mod_1"] = f"case01_rn_{suffix}"
        loader_inputs["mod_2"] = f"case01_wv_{suffix}"
        wf["40"] = {"class_type": "MiniMaxH3RefModsLoader", "inputs": loader_inputs}
        wf["41"] = {"class_type": "MiniMaxH3RefModApply", "inputs": {
            "conditioning": ["30", 0], "mods": ["40", 0],
            "override": False, "retention": 1.0, "curve_direction": "constant",
            "curve_shape": "linear", "curve_value": 1.0,
            "scramble_seed": -1, "max_total_tokens": 0}}
        pos = ["41", 0]
        latent_src = ["30", 1]

    # negative cond: text-only ReferenceToVideo without refs
    wf["31"] = {"class_type": "MiniMaxH3ReferenceToVideo", "inputs": {
        "clip": ["10", 0], "vae": ["11", 0], "audio_vae": ["12", 0],
        "prompt": NEG, "width": WIDTH, "height": HEIGHT, "length": LENGTH,
        "ref_image_size": "match"}}
    # sampler chain (native two: sigma shift inside? use plain KSampler w/ Euler+simple; ref2va standard = res_multistep 20)
    wf["50"] = {"class_type": "KSamplerSelect", "inputs": {"sampler_name": "euler"}}
    wf["51"] = {"class_type": "BasicScheduler", "inputs": {"model": ["13", 0], "scheduler": "simple", "steps": 20, "denoise": 1.0}}
    wf["52"] = {"class_type": "BasicGuider", "inputs": {"model": ["13", 0], "conditioning": pos}}
    wf["53"] = {"class_type": "RandomNoise", "inputs": {"noise_seed": SEED, "control": "fixed"}}
    wf["54"] = {"class_type": "SamplerCustomAdvanced", "inputs": {
        "noise": ["53", 0], "guider": ["52", 0], "sampler": ["50", 0], "sigmas": ["51", 0], "latent_image": latent_src}}
    wf["60"] = {"class_type": "VAEDecode", "inputs": {"samples": ["54", 0], "vae": ["11", 0]}}
    
    return wf

def find_save_node_cls():
    oi = json.loads(opener.open(BASE + "/object_info/SaveVIDEO", timeout=15).read())
    return "SaveVIDEO" in oi

if __name__ == "__main__":
    arms = sys.argv[1:] or ["N", "E", "T"]
    # resolve actual save node name
    oi = json.loads(opener.open(BASE + "/object_info", timeout=60).read())
    save_candidates = [k for k in oi if k.upper() in ("SAVEVIDEO", "SAVEWEBM", "VHS_VIDEORECORD")]
    print("save node candidates:", save_candidates, flush=True)
    for arm in arms:
        wf = build_arm(arm)
        # adapt save node
        sc = save_candidates[0]
        if sc == "SaveVIDEO":
            wf["61"] = {"class_type": sc, "inputs": {"images": ["60", 0], "filename_prefix": f"refmod_probe_arm{arm}", "format": "video/h264-mp4", "fps": 24.0}}
        elif sc == "SaveWEBM":
            wf["61"] = {"class_type": sc, "inputs": {"images": ["60", 0], "filename_prefix": f"refmod_probe_arm{arm}", "codec": "vp9", "fps": 24.0, "crf": 20.0}}
        else:
            wf["61"] = {"class_type": sc, "inputs": {"images": ["60", 0], "filename_prefix": f"refmod_probe_arm{arm}"}}
        # check clip name exists
        clip_opts = oi["CLIPLoader"]["input"]["required"]["clip_name"][0]
        clips = [c for c in clip_opts if "qwen" in c.lower()]
        print("clip options with qwen:", clips[:3], flush=True)
        wf["10"]["inputs"]["clip_name"] = clips[0]
        t0 = time.time()
        try:
            pid = post(wf)
            print(f"[submit] arm {arm} -> {pid}", flush=True)
            status, errs, outs = wait(pid)
            dt = time.time() - t0
            errinfo = json.dumps(errs[0][1])[:400] if errs else ""
            print(f"[done  ] arm {arm}: {status} in {dt:.0f}s {errinfo}", flush=True)
            if outs: print("  outputs:", json.dumps(outs)[:200], flush=True)
        except Exception as e:
            print(f"[FAIL  ] arm {arm}: {e}", flush=True)
    print("ALL_DONE")