#!/usr/bin/env python3
"""
DolinaSmrti Asset Generator - ComfyUI API Version
Calls local ComfyUI server at http://127.0.0.1:8188
"""

import os
import json
import time
import base64
import requests
from pathlib import Path
from typing import List, Dict

# Configuration
COMFYUI_URL = "http://127.0.0.1:8000"
OUTPUT_DIR = "/Users/davidkotnik/repos/novafarma/assets/images"
STYLE_PREFIX = "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, "

# Master Registry v1.14 - COMPLETE ASSET LIST
ASSETS = [
    # === CHARACTERS & NPCs ===
    {"cat": "npcs", "file": "kai_protagonist.png", "prompt": STYLE_PREFIX + "Kai protagonist with bright pink dreadlocks, athletic build, survivor clothing, determined expression, 4-direction sprite sheet"},
    {"cat": "npcs", "file": "grok_alpha_troll.png", "prompt": STYLE_PREFIX + "Grok Alpha Troll, GREEN SKIN, rainbow dreadlocks, facial piercings, stretched ear gauges, graphic t-shirt, baggy pants, holding vape device"},
    {"cat": "npcs", "file": "ana_twin.png", "prompt": STYLE_PREFIX + "Ana, Kai's twin sister, pink dreadlocks, kind expression, explorer outfit, 4-direction sprite"},
    {"cat": "npcs", "file": "ivan_blacksmith.png", "prompt": STYLE_PREFIX + "Ivan the Blacksmith, muscular build, leather apron, hammer, friendly beard, town NPC"},
    {"cat": "npcs", "file": "marija_baker.png", "prompt": STYLE_PREFIX + "Marija the Baker, warm smile, chef's apron, holding bread, town NPC"},
    
    # === LIVESTOCK - Basic ===
    {"cat": "zivali", "file": "susi_dachshund.png", "prompt": STYLE_PREFIX + "Susi the Dachshund, brown with darker spots, long body, short legs, floppy ears, 4-direction sprite"},
    {"cat": "zivali", "file": "cow_normal.png", "prompt": STYLE_PREFIX + "Normal dairy cow, black and white Holstein pattern, friendly face, 4-direction sprite"},
    {"cat": "zivali", "file": "sheep_normal.png", "prompt": STYLE_PREFIX + "Normal farm sheep, fluffy white wool, black face and legs, 4-direction sprite"},
    {"cat": "zivali", "file": "chicken_normal.png", "prompt": STYLE_PREFIX + "Normal farm chicken, white feathers, red comb, pecking stance, 4-direction sprite"},
    {"cat": "zivali", "file": "pig_normal.png", "prompt": STYLE_PREFIX + "Normal farm pig, pink colored, curly tail, round body, 4-direction sprite"},
    {"cat": "zivali", "file": "horse_normal.png", "prompt": STYLE_PREFIX + "Normal horse for riding, brown coat, black mane and tail, 4-direction sprite"},
    
    # === LIVESTOCK - Mutants ===
    {"cat": "zivali", "file": "sheep_fire.png", "prompt": STYLE_PREFIX + "Fire Sheep mutant, wool made of flickering orange-red flames, glowing effect, dark legs, magical"},
    {"cat": "zivali", "file": "sheep_golden_fleece.png", "prompt": STYLE_PREFIX + "Golden Fleece Sheep, shimmering gold metallic wool, sparkles, legendary rare"},
    {"cat": "zivali", "file": "cow_mutant.png", "prompt": STYLE_PREFIX + "Mutant Cow, larger size, glowing eyes, produces more milk, enhanced features"},
    {"cat": "zivali", "file": "chicken_three_headed.png", "prompt": STYLE_PREFIX + "Three-Headed Chicken, white body with THREE chicken heads on long necks, bizarre yet cute"},
    {"cat": "zivali", "file": "chicken_rainbow.png", "prompt": STYLE_PREFIX + "Rainbow Chicken, feathers in rainbow gradient colors, sparkle effects, magical"},
    {"cat": "zivali", "file": "chicken_phoenix.png", "prompt": STYLE_PREFIX + "Phoenix Chicken, bright orange-red feathers, flame patterns, fiery tail, glowing eyes"},
    {"cat": "zivali", "file": "pig_giant.png", "prompt": STYLE_PREFIX + "Giant Pig rideable, HUGE pink pig twice normal size, saddle on back, strong build"},
    {"cat": "zivali", "file": "horse_undead.png", "prompt": STYLE_PREFIX + "Undead Horse, dark skeletal with glowing blue ethereal mane, visible ribcage, ghostly aura"},
    {"cat": "zivali", "file": "unicorn.png", "prompt": STYLE_PREFIX + "Unicorn, pure white with rainbow gradient mane and tail, golden spiral horn, magical sparkles"},
    {"cat": "zivali", "file": "golden_goose.png", "prompt": STYLE_PREFIX + "Golden Goose, entire body shimmering gold, lays golden eggs, legendary rare"},
    
    # === DOG BREEDS ===
    {"cat": "zivali", "file": "dog_retriever.png", "prompt": STYLE_PREFIX + "Golden Retriever dog, friendly loyal companion, finds items, 4-direction sprite"},
    {"cat": "zivali", "file": "dog_shepherd.png", "prompt": STYLE_PREFIX + "German Shepherd dog, strong combat bonus, protective stance, 4-direction sprite"},
    {"cat": "zivali", "file": "dog_husky.png", "prompt": STYLE_PREFIX + "Husky dog, cold resistance, blue eyes, fluffy coat, 4-direction sprite"},
    {"cat": "zivali", "file": "dog_corgi.png", "prompt": STYLE_PREFIX + "Corgi dog, short legs, fast loyalty gain, cute expression, 4-direction sprite"},
    {"cat": "zivali", "file": "dog_dalmatian.png", "prompt": STYLE_PREFIX + "Dalmatian dog, white with black spots, speed boost ability, 4-direction sprite"},
    
    # === SLIMES (8 Types) ===
    {"cat": "mutanti", "file": "slime_green.png", "prompt": STYLE_PREFIX + "Green Slime, bouncing blob, basic slime, cute round shape"},
    {"cat": "mutanti", "file": "slime_blue.png", "prompt": STYLE_PREFIX + "Blue Slime, water element, bouncing blob, aqua colored"},
    {"cat": "mutanti", "file": "slime_red.png", "prompt": STYLE_PREFIX + "Red Slime, fire element, bouncing blob, bright red"},
    {"cat": "mutanti", "file": "slime_yellow.png", "prompt": STYLE_PREFIX + "Yellow Slime, electric element, bouncing blob, bright yellow sparks"},
    {"cat": "mutanti", "file": "slime_purple.png", "prompt": STYLE_PREFIX + "Purple Slime, poison element, bouncing blob, toxic purple"},
    {"cat": "mutanti", "file": "slime_black.png", "prompt": STYLE_PREFIX + "Black Slime, shadow element, bouncing blob, dark mysterious"},
    {"cat": "mutanti", "file": "slime_rainbow.png", "prompt": STYLE_PREFIX + "Rainbow Slime, all elements, bouncing blob, rainbow gradient shimmer"},
    {"cat": "mutanti", "file": "slime_king.png", "prompt": STYLE_PREFIX + "King Slime BOSS, GIANT blue slime with crown, menacing yet cute, royal"},
    
    # === FOREST WILDLIFE ===
    {"cat": "zivali", "file": "fox_wild.png", "prompt": STYLE_PREFIX + "Wild Fox, orange fur with white chest, bushy tail, forest animal, 4-direction sprite"},
    {"cat": "zivali", "file": "deer_wild.png", "prompt": STYLE_PREFIX + "Wild Deer, brown coat, antlers on male, peaceful forest animal, 4-direction sprite"},
    {"cat": "zivali", "file": "rabbit_wild.png", "prompt": STYLE_PREFIX + "Wild Rabbit, brown fur, long ears, hopping stance, forest animal"},
    {"cat": "zivali", "file": "hedgehog_wild.png", "prompt": STYLE_PREFIX + "Hedgehog, brown with spiky back, small cute forest animal"},
    {"cat": "zivali", "file": "bear_wild.png", "prompt": STYLE_PREFIX + "Wild Bear, large brown bear, dangerous predator, forest animal, 4-direction sprite"},
    {"cat": "zivali", "file": "wolf_wild.png", "prompt": STYLE_PREFIX + "Wild Wolf, grey fur, pack predator, dangerous forest animal, 4-direction sprite"},
    {"cat": "zivali", "file": "boar_wild.png", "prompt": STYLE_PREFIX + "Wild Boar, dark bristly fur, tusks, aggressive forest animal, 4-direction sprite"},
    {"cat": "zivali", "file": "owl_nocturnal.png", "prompt": STYLE_PREFIX + "Owl, brown feathers, wise eyes, nocturnal forest bird, perched"},
    {"cat": "zivali", "file": "bat_nocturnal.png", "prompt": STYLE_PREFIX + "Bat, flying nocturnal creature, spread wings, forest animal"},
    
    # === MARINE WILDLIFE ===
    {"cat": "zivali", "file": "fish_bass.png", "prompt": STYLE_PREFIX + "Bass fish, common freshwater fish, silver scales, swimming"},
    {"cat": "zivali", "file": "fish_trout.png", "prompt": STYLE_PREFIX + "Trout fish, spotted pattern, freshwater fish, swimming"},
    {"cat": "zivali", "file": "fish_salmon.png", "prompt": STYLE_PREFIX + "Salmon fish, orange-pink flesh color, ocean fish, swimming"},
    {"cat": "zivali", "file": "fish_tuna.png", "prompt": STYLE_PREFIX + "Tuna fish, large ocean fish, blue-grey coloring, swimming"},
    {"cat": "zivali", "file": "fish_golden.png", "prompt": STYLE_PREFIX + "Golden Fish, legendary rare shimmering gold fish, magical sparkles"},
    {"cat": "zivali", "file": "piranha.png", "prompt": STYLE_PREFIX + "Piranha, dangerous Amazon fish, sharp teeth, aggressive"},
    {"cat": "zivali", "file": "shark.png", "prompt": STYLE_PREFIX + "Shark, ocean predator, grey with fin, dangerous marine animal"},
    {"cat": "zivali", "file": "jellyfish.png", "prompt": STYLE_PREFIX + "Jellyfish, translucent bell shape, trailing tentacles, ocean creature"},
]

def queue_comfy_prompt(prompt_text: str, output_name: str) -> dict:
    """
    Send prompt to ComfyUI and queue generation
    Returns: workflow ID
    """
    # Basic ComfyUI API workflow
    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 queuing prompt: {e}")
        return None


def wait_for_completion(prompt_id: str, timeout: int = 180) -> 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:
                status = history[prompt_id].get("status", {})
                if status.get("completed", False):
                    return True
                    
        except Exception as e:
            print(f"⚠️ Polling error: {e}")
        
        time.sleep(2)
    
    return False


def download_image(prompt_id: str, output_path: Path) -> bool:
    """Download generated image from ComfyUI"""
    try:
        # Get image filename from history
        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
                    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


def main():
    """Main generation loop"""
    print("=" * 70)
    print("🎨 DolinaSmrti Asset Generator - ComfyUI Local")
    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")
        print("✅ ComfyUI server is running!")
    except:
        print("❌ ERROR: ComfyUI server not responding at http://127.0.0.1:8188")
        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"   📝 Category: {asset['cat']}")
        
        # Queue prompt
        result = queue_comfy_prompt(asset["prompt"], asset["file"].replace(".png", ""))
        if not result or "prompt_id" not in result:
            print(f"   ❌ FAILED to queue")
            fail_count += 1
            continue
        
        prompt_id = result["prompt_id"]
        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
        
        # Report every 10 images
        if (i) % 10 == 0:
            print(f"\n📊 Progress: {success_count} success, {skip_count} skipped, {fail_count} failed\n")
        
        time.sleep(1)  # Small delay
    
    # 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()