#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Generate H3 T2V test: API prompt JSON + LiteGraph workflow JSON (with dasiwa MiniMax H3 Cache)."""
import json, sys

PROMPT = """integrated_multimodal_description: A cinematic drone shot gliding low over a misty mountain lake at sunrise. Golden light breaks through layered fog, revealing dark pine forests on steep slopes. Gentle ripples catch the warm glow, thin clouds drift across the peaks. The camera moves forward and slightly upward, revealing more of the valley. Atmospheric, serene, ultra detailed.
overall_soundscape: Gentle morning wind over water, distant bird calls, soft rustling of pine trees.
non_diegetic_music: Subtle ambient orchestral pad, warm and uplifting, slow build."""

W, H, LEN = 640, 640, 243  # 0.41 MP, 10.1 s @24fps

# ---------------- API prompt format ----------------
def P(class_type, **inputs):
    return {"class_type": class_type, "inputs": inputs}

prompt = {
    "1": P("UNETLoader", unet_name="minimax_h3_fl2va_pruned_w4a8_mixed.safetensors", weight_dtype="default"),
    "2": P("LoraLoaderModelOnly", model=["1", 0], lora_name="minimax/minimax_h3_fl2v_turbo_8step_v1.0_comfyui_bf16.safetensors", strength_model=1.0),
    "3": P("MiniMaxH3Cache", model=["2", 0], reuse_threshold=0.05, start_percent=0.15, end_percent=0.90, max_steps=2, device="auto", verbose=True),
    "4": P("PathchComfyKitchenAttentionDaSiWa", model=["3", 0]),
    "5": P("CLIPLoader", clip_name="qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", type="minimax", device="default"),
    "6": P("VAELoader", vae_name="minimax_h3_video_vae_int8_convrot.safetensors"),
    "7": P("VAELoader", vae_name="minimax_h3_audio_vae_fp32.safetensors"),
    "8": P("MiniMaxH3ImageToVideo", clip=["5", 0], vae=["6", 0], prompt=PROMPT, width=W, height=H, length=LEN),
    "9": P("RandomNoise", noise_seed=20260829),
    "10": P("BasicGuider", model=["4", 0], conditioning=["8", 0]),
    "11": P("KSamplerSelect", sampler_name="res_multistep"),
    "12": P("BasicScheduler", model=["4", 0], scheduler="simple", steps=8, denoise=1.0),
    "13": P("SamplerCustomAdvanced", noise=["9", 0], guider=["10", 0], sampler=["11", 0], sigmas=["12", 0], latent_image=["8", 1]),
    "14": P("VAEDecode", samples=["13", 0], vae=["6", 0]),
    "15": P("VAEDecodeAudio", samples=["13", 0], vae=["7", 0]),
    "16": P("CreateVideo", images=["14", 0], audio=["15", 0], fps=24, bit_depth=8),
    "17": P("SaveVideo", video=["16", 0], filename_prefix="video/MiniMax_H3_cache", format="auto", codec="auto"),
}

with open("h3_t2v_cache_prompt.json", "w") as f:
    json.dump({"prompt": prompt, "client_id": "dasiwa-h3-test"}, f, ensure_ascii=False, indent=1)

# ---------------- LiteGraph workflow format ----------------
def slot(name, stype, links=None):
    s = {"name": name, "type": stype, "links": links or []}
    return s

def wf_node(nid, ntype, x, y, widgets, inputs, outputs, w=280, h=120):
    return {
        "id": nid, "type": ntype, "pos": [x, y], "size": [w, h], "flags": {},
        "order": nid, "mode": 0, "inputs": inputs, "outputs": outputs,
        "properties": {"Node name for S&R": ntype}, "widgets_values": widgets,
    }

L = []  # links: [id, from_id, from_slot, to_id, to_slot, type]
lid = 1
def link(fid, fslot, tid, tslot, ttype):
    global lid
    L.append([lid, fid, fslot, tid, tslot, ttype]); lid += 1

nodes = []
def add(n):
    nodes.append(n); return n["id"]

x0, y0, dx, dy = -1200, -600, 320, 150
def pos(i, col):
    return [x0 + col*dx, y0 + i*dy]

u1 = add(wf_node(1, "UNETLoader", *pos(0,0), ["minimax_h3_fl2va_pruned_w4a8_mixed.safetensors","default"],
                [slot("unet_name","STRING"), slot("weight_dtype","STRING")], [slot("MODEL","MODEL")]))
u2 = add(wf_node(2, "LoraLoaderModelOnly", *pos(1,0), ["minimax/minimax_h3_fl2v_turbo_8step_v1.0_comfyui_bf16.safetensors", 1.0],
                [slot("model","MODEL"), slot("lora_name","STRING"), slot("strength_model","FLOAT")], [slot("MODEL","MODEL")]))
u3 = add(wf_node(3, "MiniMaxH3Cache", *pos(2,0), [0.05, 0.15, 0.90, 2, "auto", True],
                [slot("model","MODEL"), slot("reuse_threshold","FLOAT"), slot("start_percent","FLOAT"), slot("end_percent","FLOAT"), slot("max_steps","INT"), slot("device",["auto","cuda","cpu"]), slot("verbose","BOOLEAN")],
                [slot("MODEL","MODEL")]))
u4 = add(wf_node(4, "PathchComfyKitchenAttentionDaSiWa", *pos(3,0), [],
                [slot("model","MODEL")], [slot("MODEL","MODEL")]))
u5 = add(wf_node(5, "CLIPLoader", *pos(0,2), ["qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors","minimax","default"],
                [slot("clip_name","STRING"), slot("type","STRING"), slot("device","STRING")], [slot("CLIP","CLIP")]))
u6 = add(wf_node(6, "VAELoader", *pos(1,2), ["minimax_h3_video_vae_int8_convrot.safetensors"],
                [slot("vae_name","STRING")], [slot("VAE","VAE")]))
u7 = add(wf_node(7, "VAELoader", *pos(2,2), ["minimax_h3_audio_vae_fp32.safetensors"],
                [slot("vae_name","STRING")], [slot("VAE","VAE")]))
u8 = add(wf_node(8, "MiniMaxH3ImageToVideo", *pos(3,2), [PROMPT, W, H, LEN],
                [slot("clip","CLIP"), slot("vae","VAE"), slot("prompt","STRING"), slot("width","INT"), slot("height","INT"), slot("length","INT"), slot("first_frame","IMAGE"), slot("last_frame","IMAGE")],
                [slot("positive","CONDITIONING"), slot("LATENT","LATENT")]))
u9 = add(wf_node(9, "RandomNoise", *pos(0,4), [20260829, "fixed"],
                [slot("noise_seed","INT")], [slot("NOISE","NOISE")]))
u10 = add(wf_node(10, "BasicGuider", *pos(1,4), [],
                [slot("model","MODEL"), slot("conditioning","CONDITIONING")], [slot("GUIDER","GUIDER")]))
u11 = add(wf_node(11, "KSamplerSelect", *pos(2,4), ["res_multistep"],
                [slot("sampler_name","STRING")], [slot("SAMPLER","SAMPLER")]))
u12 = add(wf_node(12, "BasicScheduler", *pos(3,4), ["simple", 8, 1.0],
                [slot("model","MODEL"), slot("scheduler","STRING"), slot("steps","INT"), slot("denoise","FLOAT")], [slot("SIGMAS","SIGMAS")]))
u13 = add(wf_node(13, "SamplerCustomAdvanced", *pos(0,6), [],
                [slot("noise","NOISE"), slot("guider","GUIDER"), slot("sampler","SAMPLER"), slot("sigmas","SIGMAS"), slot("latent_image","LATENT")],
                [slot("output","LATENT"), slot("denoised_output","LATENT")]))
u14 = add(wf_node(14, "VAEDecode", *pos(1,6), [],
                [slot("samples","LATENT"), slot("vae","VAE")], [slot("IMAGE","IMAGE")]))
u15 = add(wf_node(15, "VAEDecodeAudio", *pos(2,6), [],
                [slot("samples","LATENT"), slot("vae","VAE")], [slot("AUDIO","AUDIO")]))
u16 = add(wf_node(16, "CreateVideo", *pos(3,6), [24, 8],
                [slot("images","IMAGE"), slot("audio","AUDIO"), slot("fps","INT"), slot("bit_depth","INT")], [slot("VIDEO","VIDEO")]))
u17 = add(wf_node(17, "SaveVideo", *pos(3,8), ["video/MiniMax_H3_cache", "auto", "auto"],
                [slot("video","VIDEO"), slot("filename_prefix","STRING"), slot("format","STRING"), slot("codec","STRING")], [slot("video","VIDEO")]))

link(u1,0,u2,0,"MODEL"); link(u2,0,u3,0,"MODEL"); link(u3,0,u4,0,"MODEL")
link(u4,0,u10,0,"MODEL"); link(u4,0,u12,0,"MODEL")
link(u5,0,u8,0,"CLIP"); link(u6,0,u8,1,"VAE")
link(u8,0,u10,1,"CONDITIONING"); link(u8,1,u13,4,"LATENT")
link(u9,0,u13,0,"NOISE"); link(u10,0,u13,1,"GUIDER"); link(u11,0,u13,2,"SAMPLER"); link(u12,0,u13,3,"SIGMAS")
link(u13,0,u14,0,"LATENT"); link(u13,0,u15,0,"LATENT")
link(u6,0,u14,1,"VAE"); link(u7,0,u15,1,"VAE")
link(u14,0,u16,0,"IMAGE"); link(u15,0,u16,1,"AUDIO")
link(u16,0,u17,0,"VIDEO")

workflow = {
    "last_node_id": 17, "last_link_id": lid-1, "nodes": nodes, "links": L,
    "groups": [], "config": {}, "extra": {}, "version": 0.4,
}
with open("h3_t2v_cache_workflow.json", "w") as f:
    json.dump(workflow, f, ensure_ascii=False)

print("prompt nodes:", len(prompt), "| workflow nodes:", len(nodes), "links:", len(L))