#!/usr/bin/env python3
"""Verify the user's workflow AS-IS (boogu_image_base_bf16 + flux1_vae_bf16 + boonude LoRA,
official-edit-style sampling: AuraFlow 3.16, dpmpp_2m, simple/25/1, SamplerCustom cfg 3.5).
Core generation path only; safe prompt for mechanics test (workflow file itself untouched)."""
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)


GRAPH = {
    "1": {"class_type": "UNETLoader", "inputs": {"unet_name": "boogu_image_base_bf16.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": "flux1_vae_bf16.safetensors"}},
    "4": {"class_type": "LoraLoaderModelOnly", "inputs": {"model": ["1", 0], "lora_name": "boonude.safetensors", "strength_model": 1.0}},
    "5": {"class_type": "ModelSamplingAuraFlow", "inputs": {"model": ["1", 0], "shift": 3.16}},
    "6": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 0], "text": "A beautiful portrait of a woman in a sunlit garden, soft cinematic lighting, professional photography"}},
    "7": {"class_type": "ConditioningZeroOut", "inputs": {"conditioning": ["6", 0]}},
    "8": {"class_type": "EmptyLatentImage", "inputs": {"width": 1024, "height": 1024, "batch_size": 1}},
    "9": {"class_type": "KSamplerSelect", "inputs": {"sampler_name": "dpmpp_2m"}},
    "10": {"class_type": "BasicScheduler", "inputs": {"scheduler": "simple", "steps": 25, "denoise": 1.0, "model": ["5", 0]}},
    "11": {"class_type": "SamplerCustom", "inputs": {"model": ["4", 0], "add_noise": True, "noise_seed": 22, "cfg": 3.5, "positive": ["6", 0], "negative": ["7", 0], "sampler": ["9", 0], "sigmas": ["10", 0], "latent_image": ["8", 0]}},
    "12": {"class_type": "VAEDecode", "inputs": {"samples": ["11", 0], "vae": ["3", 0]}},
    "13": {"class_type": "SaveImage", "inputs": {"images": ["12", 0], "filename_prefix": "Boogu_base_bf16_verify"}},
}

try:
    resp = api("/prompt", {"prompt": GRAPH, "client_id": "boogu-bf16-verify"})
except Exception as e:
    print("submit error:", e)
    sys.exit(1)
if "prompt_id" not in resp:
    print("validation failed:", json.dumps(resp, ensure_ascii=False)[:1500])
    sys.exit(2)
pid = resp["prompt_id"]
print("submitted", 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"status={st} elapsed={time.time()-t0:.0f}s", flush=True)
        for m in p.get("status", {}).get("messages", []):
            if m[0] == "execution_error":
                print("ERROR:", json.dumps(m[1], ensure_ascii=False)[:1500])
        for nid, out in p.get("outputs", {}).items():
            for im in out.get("images", []):
                print("IMAGE:", im["filename"])
        sys.exit(0 if st == "success" else 3)