#!/usr/bin/env python3
"""
Test script - generates single asset to verify ComfyUI workflow
"""
import os
import json
import time
import requests
from pathlib import Path

COMFYUI_URL = "http://127.0.0.1:8000"
OUTPUT_DIR = "/Users/davidkotnik/repos/novafarma/assets/images"

# Test prompt with all requirements
TEST_PROMPT = "Top-down 2.5D video game sprite, 128x128 pixels, smooth vector art style like Stardew Valley, NO pixels, vibrant colors, clean edges, GREEN SCREEN background RGBA(0,255,0,255) pure green chroma key, Normal farm cow black and white Holstein pattern friendly face 4-direction sprite"

def queue_comfy_prompt(prompt_text: str, output_name: str) -> dict:
    """Send prompt to ComfyUI"""
    workflow = {
        "3": {
            "inputs": {
                "seed": int(time.time()),
                "steps": 25,
                "cfg": 7.5,
                "sampler_name": "euler",
                "scheduler": "normal",
                "denoise": 1,
                "model": ["4", 0],
                "positive": ["6", 0],
                "negative": ["7", 0],
                "latent_image": ["5", 0]
            },
            "class_type": "KSampler"
        },
        "4": {
            "inputs": {
                "ckpt_name": "sd_xl_base_1.0.safetensors"
            },
            "class_type": "CheckpointLoaderSimple"
        },
        "5": {
            "inputs": {
                "width": 128,
                "height": 128,
                "batch_size": 1
            },
            "class_type": "EmptyLatentImage"
        },
        "6": {
            "inputs": {
                "text": prompt_text,
                "clip": ["4", 1]
            },
            "class_type": "CLIPTextEncode"
        },
        "7": {
            "inputs": {
                "text": "blurry, low quality, pixelated, voxel, 3D render, realistic photo",
                "clip": ["4", 1]
            },
            "class_type": "CLIPTextEncode"
        },
        "8": {
            "inputs": {
                "samples": ["3", 0],
                "vae": ["4", 2]
            },
            "class_type": "VAEDecode"
        },
        "9": {
            "inputs": {
                "filename_prefix": output_name,
                "images": ["8", 0]
            },
            "class_type": "SaveImage"
        }
    }
    
    try:
        response = requests.post(
            f"{COMFYUI_URL}/prompt",
            json={"prompt": workflow}
        )
        return response.json()
    except Exception as e:
        print(f"❌ Error: {e}")
        return None

def wait_for_completion(prompt_id: str, timeout: int = 180) -> bool:
    """Wait for generation"""
    start = time.time()
    while time.time() - start < timeout:
        try:
            response = requests.get(f"{COMFYUI_URL}/history/{prompt_id}")
            history = response.json()
            
            if prompt_id in history:
                status = history[prompt_id].get("status", {})
                if status.get("completed", False):
                    return True
                    
        except Exception as e:
            print(f"⚠️ Polling: {e}")
        
        time.sleep(2)
    
    return False

def download_image(prompt_id: str, output_path: Path) -> bool:
    """Download image"""
    try:
        response = requests.get(f"{COMFYUI_URL}/history/{prompt_id}")
        history = response.json()
        
        if prompt_id not in history:
            return False
            
        outputs = history[prompt_id].get("outputs", {})
        for node_id, node_output in outputs.items():
            if "images" in node_output:
                for img in node_output["images"]:
                    filename = img["filename"]
                    subfolder = img.get("subfolder", "")
                    
                    img_url = f"{COMFYUI_URL}/view"
                    params = {
                        "filename": filename,
                        "subfolder": subfolder,
                        "type": "output"
                    }
                    
                    img_response = requests.get(img_url, params=params)
                    if img_response.status_code == 200:
                        output_path.parent.mkdir(parents=True, exist_ok=True)
                        with open(output_path, 'wb') as f:
                            f.write(img_response.content)
                        return True
        
        return False
        
    except Exception as e:
        print(f"❌ Download error: {e}")
        return False

print("=" * 70)
print("🧪 TEST: Single Asset Generation")
print("=" * 70)
print(f"📡 Server: {COMFYUI_URL}")
print(f"🎨 Prompt: {TEST_PROMPT[:80]}...")
print()

# Check connection
try:
    response = requests.get(f"{COMFYUI_URL}/system_stats")
    print("✅ ComfyUI connected!")
except:
    print("❌ ComfyUI NOT responding!")
    exit(1)

# Queue test
print("\n🚀 Queuing test generation...")
result = queue_comfy_prompt(TEST_PROMPT, "test_cow")
if not result or "prompt_id" not in result:
    print("❌ Failed to queue!")
    exit(1)

prompt_id = result["prompt_id"]
print(f"⏳ Queued: {prompt_id}")

# Wait
print("⏳ Waiting for completion...")
if wait_for_completion(prompt_id):
    output_path = Path(OUTPUT_DIR) / "zivali" / "test_cow.png"
    if download_image(prompt_id, output_path):
        print(f"✅ SUCCESS! Saved to: {output_path}")
        print("\n📊 Check the image:")
        print(f"   - Size should be 128x128")
        print(f"   - Background should be green RGBA(0,255,0,255)")
        print(f"   - Style should be smooth 2D vector art")
    else:
        print("❌ Failed to download!")
else:
    print("⏱️ Timeout!")