"""Submit and time a prompt-only MiniMax H3 cat benchmark in ComfyUI."""

from __future__ import annotations

import argparse
import json
import time
from pathlib import Path
from urllib.request import Request, urlopen


DEFAULT_PROMPT = (
    "A playful orange tabby cat starts in an ordinary cozy living room in a normal house, afternoon light, sofa and rug. "
    "The cat crouches, jumps, and does one clean athletic backflip in slow motion. "
    "As the backflip completes there is a sharp cinematic cut: the cat lands perfectly on a glowing neon disco dance floor wearing oversized black sunglasses. "
    "Mirror ball reflections, colorful lights, joyful party energy, stylish and funny, clear before-and-after transformation."
)


def post_json(url: str, payload: dict) -> dict:
    request = Request(url, data=json.dumps(payload).encode(), headers={"Content-Type": "application/json"})
    with urlopen(request, timeout=30) as response:
        return json.loads(response.read().decode())


def get_json(url: str) -> dict:
    with urlopen(url, timeout=30) as response:
        return json.loads(response.read().decode())


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--server", default="http://127.0.0.1:8188")
    parser.add_argument("--prompt", default=DEFAULT_PROMPT)
    parser.add_argument("--seed", type=int, default=440407)
    parser.add_argument("--width", type=int, default=960)
    parser.add_argument("--height", type=int, default=544)
    parser.add_argument("--frames", type=int, default=124)
    parser.add_argument("--steps", type=int, default=12)
    parser.add_argument("--filename-prefix", default="h3-blackwell-runtime/comfy-cat-benchmark-960x544-124f-seed440407")
    parser.add_argument("--metrics", type=Path, default=Path("/workspace/ComfyUI/output/h3-blackwell-runtime/benchmarks/comfy-cat-benchmark-960x544-124f-seed440407.json"))
    args = parser.parse_args()

    graph = {
        "1": {"class_type": "UNETLoader", "inputs": {"unet_name": "minimax_h3_fl2va_pruned_nvfp4.safetensors", "weight_dtype": "default"}},
        "3": {"class_type": "CLIPLoader", "inputs": {"clip_name": "qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", "type": "minimax"}},
        "4": {"class_type": "VAELoader", "inputs": {"vae_name": "minimax_h3_video_vae_fp16.safetensors"}},
        "5": {"class_type": "VAELoader", "inputs": {"vae_name": "minimax_h3_audio_vae_fp32.safetensors"}},
        "8": {"class_type": "MiniMaxH3ImageToVideo", "inputs": {"clip": ["3", 0], "vae": ["4", 0], "prompt": args.prompt, "width": args.width, "height": args.height, "length": args.frames}},
        "9": {"class_type": "BasicGuider", "inputs": {"model": ["1", 0], "conditioning": ["8", 0]}},
        "10": {"class_type": "RandomNoise", "inputs": {"noise_seed": args.seed}},
        "11": {"class_type": "KSamplerSelect", "inputs": {"sampler_name": "res_multistep"}},
        "12": {"class_type": "BasicScheduler", "inputs": {"model": ["1", 0], "scheduler": "beta", "steps": args.steps, "denoise": 1.0}},
        "13": {"class_type": "SamplerCustomAdvanced", "inputs": {"noise": ["10", 0], "guider": ["9", 0], "sampler": ["11", 0], "sigmas": ["12", 0], "latent_image": ["8", 1]}},
        "14": {"class_type": "VAEDecode", "inputs": {"samples": ["13", 0], "vae": ["4", 0]}},
        "15": {"class_type": "VAEDecodeAudio", "inputs": {"samples": ["13", 0], "vae": ["5", 0]}},
        "16": {"class_type": "CreateVideo", "inputs": {"images": ["14", 0], "audio": ["15", 0], "bit_depth": 8, "fps": 24.0}},
        "17": {"class_type": "SaveVideo", "inputs": {"video": ["16", 0], "filename_prefix": args.filename_prefix, "format": "mp4", "codec": "auto"}},
    }

    submitted_at = time.perf_counter()
    response = post_json(f"{args.server}/prompt", {"prompt": graph})
    prompt_id = response["prompt_id"]
    history = None
    while True:
        time.sleep(1.0)
        history_response = get_json(f"{args.server}/history/{prompt_id}")
        history = history_response.get(prompt_id)
        if history is not None:
            status = history.get("status", {})
            if status.get("completed") or status.get("status_str") in {"success", "error"}:
                break
    completed_at = time.perf_counter()

    metrics = {
        "prompt_id": prompt_id,
        "seed": args.seed,
        "width": args.width,
        "height": args.height,
        "frames": args.frames,
        "steps": args.steps,
        "wall_seconds": completed_at - submitted_at,
        "response": response,
        "status": history.get("status", {}),
        "outputs": history.get("outputs", {}),
    }
    args.metrics.parent.mkdir(parents=True, exist_ok=True)
    args.metrics.write_text(json.dumps(metrics, indent=2), encoding="utf-8")
    print(json.dumps(metrics, indent=2), flush=True)


if __name__ == "__main__":
    main()