#!/usr/bin/env python3
"""Verify Boogu-Image-0.1 official workflows on GPU machine via ComfyUI API."""
import json
import sys
import time
import urllib.request

BASE = "http://127.0.0.1:8189"


def api(path, data=None, timeout=120):
    if data is not None:
        req = urllib.request.Request(
            BASE + path,
            data=json.dumps(data).encode(),
            headers={"Content-Type": "application/json"},
        )
    else:
        req = urllib.request.Request(BASE + path)
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.load(r)


print("waiting for ComfyUI ...", flush=True)
up = False
for _ in range(120):
    try:
        s = api("/system_stats", timeout=10)
        print("server up, comfyui:", s.get("system", {}).get("comfyui_version", "?"), flush=True)
        up = True
        break
    except Exception:
        time.sleep(2)
if not up:
    print("FATAL: server did not come up", flush=True)
    sys.exit(1)

TURBO = {
    "1": {"class_type": "UNETLoader", "inputs": {"unet_name": "boogu_image_turbo_fp8_scaled.safetensors", "weight_dtype": "default"}},
    "2": {"class_type": "CLIPLoader", "inputs": {"clip_name": "qwen3vl_8b_fp8_scaled.safetensors", "type": "boogu", "device": "default"}},
    "3": {"class_type": "VAELoader", "inputs": {"vae_name": "ae.safetensors"}},
    "4": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 0], "text": "A red panda sitting on a mossy rock in a misty forest, cinematic lighting, highly detailed fur, professional wildlife photography"}},
    "5": {"class_type": "ConditioningZeroOut", "inputs": {"conditioning": ["4", 0]}},
    "6": {"class_type": "EmptyLatentImage", "inputs": {"width": 1024, "height": 1024, "batch_size": 1}},
    "7": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "seed": 896977722960984, "steps": 4, "cfg": 1.0, "sampler_name": "lcm", "scheduler": "sgm_uniform", "denoise": 1.0, "positive": ["4", 0], "negative": ["5", 0], "latent_image": ["6", 0]}},
    "8": {"class_type": "VAEDecode", "inputs": {"samples": ["7", 0], "vae": ["3", 0]}},
    "9": {"class_type": "SaveImage", "inputs": {"images": ["8", 0], "filename_prefix": "Boogu_verify_turbo"}},
}

EDIT = {
    "1": {"class_type": "UNETLoader", "inputs": {"unet_name": "boogu_image_edit_fp8_scaled.safetensors", "weight_dtype": "default"}},
    "2": {"class_type": "CLIPLoader", "inputs": {"clip_name": "qwen3vl_8b_fp8_scaled.safetensors", "type": "boogu", "device": "default"}},
    "3": {"class_type": "VAELoader", "inputs": {"vae_name": "ae.safetensors"}},
    "4": {"class_type": "LoadImage", "inputs": {"image": "tech_cowboy.png"}},
    "5": {"class_type": "ModelSamplingAuraFlow", "inputs": {"model": ["1", 0], "shift": 3.16}},
    "6": {"class_type": "TextEncodeBooguEdit", "inputs": {"clip": ["2", 0], "prompt": "remove the hat", "negative_prompt": "", "vae": ["3", 0], "images": {"image_1": ["4", 0]}}},
    "7": {"class_type": "KSamplerSelect", "inputs": {"sampler_name": "dpmpp_2m"}},
    "8": {"class_type": "BasicScheduler", "inputs": {"scheduler": "simple", "steps": 25, "denoise": 1.0, "model": ["5", 0]}},
    "9": {"class_type": "EmptyLatentImage", "inputs": {"width": 1024, "height": 1024, "batch_size": 1}},
    "10": {"class_type": "SamplerCustom", "inputs": {"model": ["5", 0], "add_noise": True, "noise_seed": 22, "cfg": 3.5, "positive": ["6", 0], "negative": ["6", 1], "sampler": ["7", 0], "sigmas": ["8", 0], "latent_image": ["9", 0]}},
    "11": {"class_type": "VAEDecode", "inputs": {"samples": ["10", 0], "vae": ["3", 0]}},
    "12": {"class_type": "SaveImage", "inputs": {"images": ["11", 0], "filename_prefix": "Boogu_verify_edit"}},
}


def run(label, prompt):
    try:
        resp = api("/prompt", {"prompt": prompt, "client_id": "boogu-verify"})
    except Exception as e:
        print(f"{label} submit error: {e}", flush=True)
        return "submit-error"
    if "prompt_id" not in resp:
        print(f"{label} validation failed: {json.dumps(resp, ensure_ascii=False)[:3000]}", flush=True)
        return "validation-error"
    pid = resp["prompt_id"]
    print(f"{label} submitted, pid={pid}", flush=True)
    t0 = time.time()
    while True:
        time.sleep(15)
        try:
            h = api("/history/" + pid, timeout=30)
        except Exception:
            continue
        if pid in h:
            p = h[pid]
            st = p["status"]["status_str"]
            print(f"{label} status={st} elapsed={time.time() - t0:.0f}s", flush=True)
            for m in p.get("status", {}).get("messages", []):
                if m[0] in ("execution_error", "execution_interrupted"):
                    print(f"{label} ERRORMSG: {json.dumps(m[1], ensure_ascii=False)[:3000]}", flush=True)
            for nid, out in p.get("outputs", {}).items():
                for im in out.get("images", []):
                    print(f"{label} IMAGE: {im.get('subfolder','')}/{im['filename']} (type={im['type']})", flush=True)
            return st


run("TURBO", TURBO)
run("EDIT", EDIT)
print("VERIFY_DONE", flush=True)