#!/usr/bin/env python3
"""Extract RefMods for Case01 characters (RN + WV), encode + training modes each.
Lessons baked in: dotted autogrow keys (refs_image.ref_image_N), LoadImage upstream,
RefModSave requires filename_prefix+subfolder."""
import json, urllib.request, time, sys

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

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"})
    return json.loads(opener.open(req, timeout=timeout).read())

def wait(prompt_id, timeout=900):
    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", [])
            err = [m for m in msgs if m[0] == "execution_error"]
            return st.get("status_str"), err, h[prompt_id].get("outputs", {})
        time.sleep(6)
    return "TIMEOUT", [], {}

def extract_workflow(name, mode, desc, img):
    return {
        "0": {"class_type": "LoadImage", "inputs": {"image": img}},
        "1": {"class_type": "VAELoader", "inputs": {"vae_name": "minimax_h3_video_vae_fp16.safetensors"}},
        "2": {"class_type": "MiniMaxH3RefModExtract", "inputs": {
            "name": name, "mode": mode, "concept_type": "identity",
            "background_retention": 0.0, "ref_resolution": 1024,
            "pool_h": 16, "pool_w": 16, "latent_frames": 4,
            "identity": 500, "merge": False, "motion_only": False,
            "multiplier": 1, "max_tokens": 5120, "description": desc,
            "save": True,
            "refs_image.ref_image_1": ["0", 0],
            "vae": ["1", 0],
        }},
        "3": {"class_type": "MiniMaxH3RefModSave", "inputs": {
            "mods": ["2", 0], "filename_prefix": name, "subfolder": ""}},
    }

RN_DESC = ("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")
WV_DESC = ("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")

jobs = [
    ("case01_rn_encode",  "encode",   RN_DESC, "kais_validation/01/1.jpg"),
    ("case01_rn_train16", "training", RN_DESC, "kais_validation/01/1.jpg"),
    ("case01_wv_encode",  "encode",   WV_DESC, "kais_validation/01/2.jpg"),
    ("case01_wv_train16", "training", WV_DESC, "kais_validation/01/2.jpg"),
]

if len(sys.argv) > 1:
    jobs = [j for j in jobs if j[0] in sys.argv[1:]]

for name, mode, desc, img in jobs:
    try:
        r = post(extract_workflow(name, mode, desc, img))
        pid = r.get("prompt_id")
        print(f"[submit] {name} -> {pid}", flush=True)
        status, errs, outs = wait(pid)
        errinfo = json.dumps(errs[0][1])[:300] if errs else ""
        print(f"[done  ] {name}: {status} {errinfo}", flush=True)
    except Exception as e:
        print(f"[FAIL  ] {name}: {e}", flush=True)
print("ALL_DONE")