#!/usr/bin/env python3
"""
DolinaSmrti Asset Generator - ComfyUI API Version
Calls local ComfyUI server at http://127.0.0.1:8188
Fixed with COMPLETE workflow and proper 2D Stardew Valley style settings
"""

import os
import json
import time
import uuid
import requests
from pathlib import Path

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

# Style for 2D Stardew Valley look (NO pixel art, smooth vector style)
STYLE_PREFIX = "Stardew Valley style, 2d game asset, flat color, clean vector art, smooth shading, cute cartoon, top-down view, game sprite, isolated on solid bright green background, "
NEGATIVE_PROMPT = "pixel art, pixelated, voxel, 3d render, realistic, photo, blurry, low quality, distorted, ugly, bad anatomy"

# Master Registry - Assets to generate
ASSETS = [
    # === ŽIVALI ===
    {"cat": "zivali", "file": "golden_goose.png", "prompt": "Golden Goose, entire body shimmering gold, lays golden eggs, legendary rare bird"},
    {"cat": "zivali", "file": "fish_trout.png", "prompt": "Trout fish, spotted pattern, freshwater fish, swimming"},
    {"cat": "zivali", "file": "shark.png", "prompt": "Shark, ocean predator, grey with fin, dangerous marine animal"},
    {"cat": "zivali", "file": "jellyfish.png", "prompt": "Jellyfish, translucent bell shape, trailing tentacles, ocean creature"},
    
    # === MUTANTI ===
    {"cat": "mutanti", "file": "slime_king.png", "prompt": "King Slime BOSS, GIANT blue slime with crown, menacing yet cute, royal"},
    
    # === DODATNI ASSETI ===
    {"cat": "items", "file": "gem_ruby.png", "prompt": "Ruby gemstone, red sparkling jewel, valuable treasure"},
    {"cat": "items", "file": "gem_emerald.png", "prompt": "Emerald gemstone, green sparkling jewel, valuable treasure"},
    {"cat": "items", "file": "gem_sapphire.png", "prompt": "Sapphire gemstone, blue sparkling jewel, valuable treasure"},
    {"cat": "items", "file": "potion_health.png", "prompt": "Health potion, red liquid in glass bottle with cork, healing item"},
    {"cat": "items", "file": "potion_mana.png", "prompt": "Mana potion, blue liquid in glass bottle with cork, magic item"},
    {"cat": "environment", "file": "tree_dead.png", "prompt": "Dead tree, bare branches, no leaves, spooky halloween tree"},
    {"cat": "environment", "file": "cactus.png", "prompt": "Cactus plant, green with spines, desert plant"},
    {"cat": "buildings", "file": "mailbox.png", "prompt": "Wooden mailbox, small post with letter slot, farm decoration"},
    {"cat": "buildings", "file": "birdhouse.png", "prompt": "Birdhouse, small wooden house for birds on pole, garden decoration"},
]


def create_workflow(prompt_text: str, output_name: str, seed: int = None) -> dict:
    """
    Creates a COMPLETE ComfyUI workflow with all required nodes.
    Uses DreamShaper model and optimized settings for 2D game assets.
    """
    if seed is None:
        seed = int(time.time() * 1000) % 2147483647
    
    workflow = {
        # Node 1: Load Checkpoint (Model)
        "1": {
            "class_type": "CheckpointLoaderSimple",
            "inputs": {
                "ckpt_name": "dreamshaper_8.safetensors"
            }
        },
        # Node 2: Empty Latent Image (Canvas size)
        "2": {
            "class_type": "EmptyLatentImage",
            "inputs": {
                "width": 512,
                "height": 512,
                "batch_size": 1
            }
        },
        # Node 3: CLIP Text Encode (Positive Prompt)
        "3": {
            "class_type": "CLIPTextEncode",
            "inputs": {
                "text": STYLE_PREFIX + prompt_text,
                "clip": ["1", 1]
            }
        },
        # Node 4: CLIP Text Encode (Negative Prompt)
        "4": {
            "class_type": "CLIPTextEncode",
            "inputs": {
                "text": NEGATIVE_PROMPT,
                "clip": ["1", 1]
            }
        },
        # Node 5: KSampler (Main generation)
        "5": {
            "class_type": "KSampler",
            "inputs": {
                "seed": seed,
                "steps": 30,
                "cfg": 7.5,
                "sampler_name": "euler_ancestral",
                "scheduler": "karras",
                "denoise": 1.0,
                "model": ["1", 0],
                "positive": ["3", 0],
                "negative": ["4", 0],
                "latent_image": ["2", 0]
            }
        },
        # Node 6: VAE Decode (Convert latent to image)
        "6": {
            "class_type": "VAEDecode",
            "inputs": {
                "samples": ["5", 0],
                "vae": ["1", 2]
            }
        },
        # Node 7: Save Image
        "7": {
            "class_type": "SaveImage",
            "inputs": {
                "filename_prefix": output_name,
                "images": ["6", 0]
            }
        }
    }
    
    return workflow


def queue_prompt(workflow: dict) -> str:
    """Send workflow to ComfyUI and return prompt_id"""
    try:
        payload = {
            "prompt": workflow,
            "client_id": f"antigravity_{uuid.uuid4().hex[:8]}"
        }
        response = requests.post(f"{COMFYUI_URL}/prompt", json=payload)
        result = response.json()
        return result.get("prompt_id")
    except Exception as e:
        print(f"❌ Queue error: {e}")
        return None


def wait_for_completion(prompt_id: str, timeout: int = 300) -> bool:
    """Wait for ComfyUI to finish generating"""
    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:
                outputs = history[prompt_id].get("outputs", {})
                if outputs:
                    return True
        except Exception as e:
            pass
        
        time.sleep(2)
    
    return False


def download_image(prompt_id: str, output_path: Path) -> bool:
    """Download generated image from ComfyUI output"""
    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", "")
                    
                    # Download image
                    params = {
                        "filename": filename,
                        "subfolder": subfolder,
                        "type": "output"
                    }
                    
                    img_response = requests.get(f"{COMFYUI_URL}/view", 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


def main():
    """Main generation loop"""
    print("=" * 70)
    print("🎨 DolinaSmrti Asset Generator - ComfyUI Local")
    print("   2D Stardew Valley Style (NO pixel art)")
    print("=" * 70)
    print(f"📡 ComfyUI Server: {COMFYUI_URL}")
    print(f"💾 Output Directory: {OUTPUT_DIR}")
    print(f"📊 Total Assets: {len(ASSETS)}")
    print()
    
    # Check ComfyUI connection
    try:
        response = requests.get(f"{COMFYUI_URL}/system_stats", timeout=5)
        data = response.json()
        print(f"✅ ComfyUI v{data['system']['comfyui_version']} is running!")
    except Exception as e:
        print(f"❌ ERROR: ComfyUI server not responding at {COMFYUI_URL}")
        print("   Please start ComfyUI first!")
        return
    
    print("\n🚀 Starting generation...\n")
    
    success_count = 0
    skip_count = 0
    fail_count = 0
    
    for i, asset in enumerate(ASSETS, 1):
        output_path = Path(OUTPUT_DIR) / asset["cat"] / asset["file"]
        
        # Skip if exists
        if output_path.exists():
            print(f"[{i}/{len(ASSETS)}] ⏭️  SKIP: {asset['file']} (exists)")
            skip_count += 1
            continue
        
        print(f"[{i}/{len(ASSETS)}] 🎨 Generating: {asset['file']}")
        print(f"   📝 {asset['prompt'][:50]}...")
        
        # Create and queue workflow
        workflow = create_workflow(asset["prompt"], asset["file"].replace(".png", ""))
        prompt_id = queue_prompt(workflow)
        
        if not prompt_id:
            print(f"   ❌ FAILED to queue")
            fail_count += 1
            continue
        
        print(f"   ⏳ Queued (ID: {prompt_id[:8]}...)")
        
        # Wait for completion
        if wait_for_completion(prompt_id):
            # Download image
            if download_image(prompt_id, output_path):
                print(f"   ✅ SAVED: {output_path}")
                success_count += 1
            else:
                print(f"   ❌ FAILED to download")
                fail_count += 1
        else:
            print(f"   ⏱️ TIMEOUT")
            fail_count += 1
        
        time.sleep(1)  # Small delay between generations
    
    # Final report
    print("\n" + "=" * 70)
    print("✅ GENERATION COMPLETE!")
    print("=" * 70)
    print(f"✅ Successful: {success_count}")
    print(f"⏭️  Skipped: {skip_count}")
    print(f"❌ Failed: {fail_count}")
    print(f"📁 Output: {OUTPUT_DIR}")


if __name__ == "__main__":
    main()