#!/usr/bin/env python3
"""Submit 3 FLUX schnell images for model comparison."""

import json
import urllib.request
import urllib.error
import time
import random

COMFYUI_URL = "http://localhost:8188"

# Base workflow
WORKFLOW = {
    "1": {
        "class_type": "UNETLoader",
        "inputs": {
            "unet_name": "flux1-schnell-fp8-e4m3fn.safetensors",
            "weight_dtype": "default"
        }
    },
    "2": {
        "class_type": "DualCLIPLoader",
        "inputs": {
            "clip_name1": "clip_l.safetensors",
            "clip_name2": "t5xxl_fp8_e4m3fn.safetensors",
            "type": "flux"
        }
    },
    "3": {
        "class_type": "VAELoader",
        "inputs": {
            "vae_name": "flux-vae-bf16.safetensors"
        }
    },
    "4": {
        "class_type": "CLIPTextEncode",
        "inputs": {
            "text": "PROMPT",
            "clip": ["2", 0]
        }
    },
    "5": {
        "class_type": "CLIPTextEncode",
        "inputs": {
            "text": "",
            "clip": ["2", 0]
        }
    },
    "6": {
        "class_type": "EmptyLatentImage",
        "inputs": {
            "width": 1024,
            "height": 576,
            "batch_size": 1
        }
    },
    "7": {
        "class_type": "KSampler",
        "inputs": {
            "seed": 42,
            "steps": 4,
            "cfg": 1.0,
            "sampler_name": "euler",
            "scheduler": "simple",
            "denoise": 1.0,
            "model": ["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": {
            "filename_prefix": "OUTPUT_PREFIX",
            "images": ["8", 0]
        }
    }
}

# Same prompts used for ZImage and Gemini
IMAGES = [
    {
        "name": "condition-ptsd-flux",
        "prompt": "Professional medical illustration for PTSD treatment article, calming purple and blue gradient, abstract neural pathways representing healing and peace, no cannabis imagery, clean modern healthcare aesthetic, gentle therapeutic feeling, text overlay area at bottom, 1024x576 banner format"
    },
    {
        "name": "city-miami-flux",
        "prompt": "Miami Florida cityscape banner for medical cannabis dispensary website, vibrant South Beach sunset colors, Art Deco buildings, palm trees silhouette, warm orange and pink sky, professional clean design, text overlay area, no cannabis imagery, 1024x576 banner format"
    },
    {
        "name": "blog-hero-weekend-deals-flux",
        "prompt": "Weekend Deals promotional banner, modern shopping aesthetic, orange and gold gradient background, abstract shopping bags and sale tags, clean professional design, excitement and savings theme, bold typography area, no cannabis or marijuana imagery, 1024x576 blog hero format"
    }
]


def submit_workflow(prompt_text: str, output_prefix: str) -> str:
    """Submit a workflow and return the prompt_id."""
    workflow = json.loads(json.dumps(WORKFLOW))
    workflow["4"]["inputs"]["text"] = prompt_text
    workflow["9"]["inputs"]["filename_prefix"] = f"canna-flux/{output_prefix}"
    workflow["7"]["inputs"]["seed"] = random.randint(1, 2**31)
    
    payload = {"prompt": workflow}
    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        f"{COMFYUI_URL}/prompt",
        data=data,
        headers={"Content-Type": "application/json"}
    )
    
    with urllib.request.urlopen(req) as resp:
        result = json.loads(resp.read().decode())
        return result.get("prompt_id")


def check_history(prompt_id: str) -> dict:
    """Check if a prompt has completed."""
    try:
        with urllib.request.urlopen(f"{COMFYUI_URL}/history/{prompt_id}") as resp:
            return json.loads(resp.read().decode())
    except urllib.error.HTTPError:
        return {}


def main():
    print("Submitting 3 FLUX schnell images...")
    prompt_ids = []
    
    for img in IMAGES:
        print(f"  Submitting: {img['name']}")
        pid = submit_workflow(img["prompt"], img["name"])
        prompt_ids.append((img["name"], pid))
        print(f"    prompt_id: {pid}")
        time.sleep(0.5)
    
    print("\nWaiting for completion...")
    completed = set()
    
    while len(completed) < len(prompt_ids):
        for name, pid in prompt_ids:
            if pid in completed:
                continue
            hist = check_history(pid)
            if hist and pid in hist:
                outputs = hist[pid].get("outputs", {})
                if "9" in outputs:
                    images = outputs["9"].get("images", [])
                    if images:
                        print(f"  ✓ {name}: {images[0]['filename']}")
                        completed.add(pid)
        if len(completed) < len(prompt_ids):
            time.sleep(2)
    
    print("\nAll done!")
    print("Output directory: /home/crogers2287/ComfyUI/output/canna-flux/")


if __name__ == "__main__":
    main()