#!/usr/bin/env python3
"""
🎮 DOLINA SMRTI - ComfyUI Local Generator
Transparent Background (Alpha Channel)
2D Indie Cartoon Style - Approved Gronk/Kai look
"""

import os
import json
import time
import uuid
import requests
from pathlib import Path
from PIL import Image
import io

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

# ============================================================================
# STYLE DEFINITIONS - Approved 2D Indie Cartoon
# ============================================================================

STYLE_PREFIX = """2D indie game sprite, cartoon vector style, clean smooth lines, 
full body visible from head to feet standing on ground, 
tight crop max 5 pixel padding around subject, """

STYLE_SUFFIX = """, isolated on pure white background #FFFFFF, game asset"""

NEGATIVE_PROMPT = """pixel art, pixelated, voxel, 3d render, realistic photo, 
anime, manga, chibi, floating, cut off, cropped, blurry, low quality, 
watermark, text, signature, green background, colored background"""

# ============================================================================
# MASTER REGISTRY
# ============================================================================

def generate_registry():
    """Generate complete asset registry"""
    assets = []
    
    # Main Characters
    gronk = """Gronk the troll, massive green-grey skinned troll, PINK dreadlocks, 
    stretched earlobes with large gauges, facial piercings, holding colorful vape device with pink smoke, 
    wearing LOOSE BAGGY black t-shirt with band logo, LOOSE WIDE black pants, pink sneakers, 
    peaceful zen expression"""
    
    kai = """Kai teenage boy survivor, post-apocalyptic nomad style, GREEN natural dreadlocks hair, 
    stretched earlobes with plugs, facial piercings on nose and lip, wearing dirty worn blue jacket, 
    torn jeans, combat boots, survival backpack, determined serious expression, athletic build"""
    
    ana = """Ana teenage girl explorer, PINK dreadlocks hair, stretched earlobes with small plugs, 
    wearing brown adventure vest over green shirt, blue jeans, brown hiking boots, 
    leather backpack with map, kind curious expression"""
    
    main_chars = [
        ("gronk_troll", gronk),
        ("kai_survivor", kai),
        ("ana_explorer", ana),
    ]
    
    for name, desc in main_chars:
        assets.append({"cat": "npcs", "file": f"{name}.png", "prompt": desc})
        for d in ["front", "back", "left", "right"]:
            for a in ["idle", "walk1", "walk2"]:
                assets.append({
                    "cat": "npcs",
                    "file": f"{name}_{d}_{a}.png",
                    "prompt": f"{desc}, {d} facing view, {a} animation frame"
                })
    
    # NPCs
    npcs = [
        ("npc_trader", "wasteland trader, nomadic merchant with cart, weathered face, hooded cloak"),
        ("npc_blacksmith", "nomadic blacksmith, leather apron, strong arms, hammer"),
        ("npc_healer", "herbalist healer woman, herb pouch, kind elderly, flower crown"),
        ("npc_hunter", "wasteland hunter, crossbow, animal pelts, face paint"),
        ("npc_farmer", "apocalypse farmer, patched overalls, straw hat, hoe"),
    ]
    
    for name, desc in npcs:
        assets.append({"cat": "npcs", "file": f"{name}.png", "prompt": desc})
    
    # Animals
    animals = [
        ("susi_dachshund", "Susi dachshund dog, brown and black, long body, floppy ears, pink collar"),
        ("dog_husky", "husky dog, fluffy grey white, blue eyes, friendly"),
        ("cat_tabby", "orange tabby cat, striped, lazy curious"),
        ("cow_spotted", "dairy cow, black white spots, friendly"),
        ("chicken_white", "white chicken, red comb, pecking"),
        ("phoenix_chicken", "Phoenix Chicken, fire feathers, glowing orange-red, legendary"),
        ("unicorn", "magical Unicorn, white horse, rainbow mane, golden spiral horn"),
    ]
    
    for name, desc in animals:
        assets.append({"cat": "zivali", "file": f"{name}.png", "prompt": desc})
    
    # Mutants
    mutants = [
        ("zombie_basic", "basic zombie, shambling, torn clothes, grey-green skin, undead"),
        ("zombie_runner", "fast zombie, sprinting, aggressive, rage face"),
        ("zombie_dreadlocks", "zombie with dreadlocks, brown shirt, farmer style"),
        ("slime_green", "green slime monster, bouncy blob, cute-creepy, small eyes"),
        ("slime_blue", "blue slime, water element, wobbly, transparent"),
    ]
    
    for name, desc in mutants:
        assets.append({"cat": "mutanti", "file": f"{name}.png", "prompt": desc})
    
    # Bosses
    bosses = [
        ("boss_trex", "T-Rex dinosaur boss, massive predator, huge jaws, fierce"),
        ("boss_kraken", "Kraken sea monster, giant octopus, many tentacles"),
        ("boss_zombie_king", "Zombie King, massive undead, rotting crown"),
        ("boss_slime_emperor", "Slime Emperor, gigantic rainbow slime, crown"),
        ("boss_dragon_fire", "Fire Dragon, red scales, breathing flames, wings"),
    ]
    
    for name, desc in bosses:
        assets.append({"cat": "bosses", "file": f"{name}.png", "prompt": desc})
    
    # Items
    items = [
        ("tool_axe_iron", "iron axe, chopping tool, wooden handle"),
        ("tool_pickaxe_iron", "iron pickaxe, mining tool"),
        ("weapon_sword_iron", "iron sword, combat blade"),
        ("weapon_bow_wood", "wooden bow with arrows"),
        ("wand_fire", "fire magic wand, red crystal tip, flames"),
        ("wand_ice", "ice magic wand, blue crystal, frost"),
        ("ana_bracelet", "Ana's magical bracelet, glowing, family heirloom"),
        ("gronk_vape", "Gronk's golden vape device, colorful, smoking"),
        ("potion_health", "health potion, red liquid bottle"),
        ("food_bread", "bread loaf, fresh baked"),
    ]
    
    for name, desc in items:
        assets.append({"cat": "items", "file": f"{name}.png", "prompt": desc})
    
    # Environment
    tiles = [
        ("tile_grass_green", "green grass ground tile, 32x32 seamless"),
        ("tile_dirt_brown", "brown dirt path tile, 32x32 seamless"),
        ("tile_water_shallow", "shallow blue water tile, 32x32 seamless"),
        ("tile_stone_grey", "grey stone floor tile, 32x32 seamless"),
    ]
    
    objects = [
        ("tree_oak", "oak tree, green leaves, brown trunk"),
        ("tree_dead", "dead tree, no leaves, gnarled branches"),
        ("rock_large", "large boulder, moss covered"),
        ("bush_green", "green bush, leafy"),
        ("campfire", "campfire, flames, logs"),
        ("chest_wooden", "wooden treasure chest, closed"),
    ]
    
    for name, desc in tiles + objects:
        assets.append({"cat": "environment", "file": f"{name}.png", "prompt": desc})
    
    return assets


# ============================================================================
# COMFYUI WORKFLOW - With Background Removal
# ============================================================================

def create_workflow(prompt_text: str, output_name: str, size: int = 512) -> dict:
    """Creates ComfyUI workflow with white background for easy removal"""
    seed = int(time.time() * 1000) % 2147483647
    full_prompt = STYLE_PREFIX + prompt_text + STYLE_SUFFIX
    
    return {
        "1": {
            "class_type": "CheckpointLoaderSimple",
            "inputs": {"ckpt_name": "dreamshaper_8.safetensors"}
        },
        "2": {
            "class_type": "EmptyLatentImage",
            "inputs": {"width": size, "height": size, "batch_size": 1}
        },
        "3": {
            "class_type": "CLIPTextEncode",
            "inputs": {"text": full_prompt, "clip": ["1", 1]}
        },
        "4": {
            "class_type": "CLIPTextEncode",
            "inputs": {"text": NEGATIVE_PROMPT, "clip": ["1", 1]}
        },
        "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]
            }
        },
        "6": {
            "class_type": "VAEDecode",
            "inputs": {"samples": ["5", 0], "vae": ["1", 2]}
        },
        "7": {
            "class_type": "SaveImage",
            "inputs": {"filename_prefix": output_name, "images": ["6", 0]}
        }
    }


def queue_prompt(workflow: dict) -> str:
    try:
        r = requests.post(
            f"{COMFYUI_URL}/prompt",
            json={"prompt": workflow, "client_id": f"ds_{uuid.uuid4().hex[:8]}"}
        )
        return r.json().get("prompt_id")
    except Exception as e:
        print(f"Queue error: {e}")
        return None


def wait_completion(prompt_id: str, timeout: int = 300) -> bool:
    start = time.time()
    while time.time() - start < timeout:
        try:
            r = requests.get(f"{COMFYUI_URL}/history/{prompt_id}")
            data = r.json()
            if prompt_id in data and data[prompt_id].get("outputs"):
                return True
        except:
            pass
        time.sleep(2)
    return False


def download_image(prompt_id: str, output_path: Path) -> bool:
    try:
        h = requests.get(f"{COMFYUI_URL}/history/{prompt_id}").json()
        for out in h.get(prompt_id, {}).get("outputs", {}).values():
            for img in out.get("images", []):
                r = requests.get(f"{COMFYUI_URL}/view", params={
                    "filename": img["filename"],
                    "subfolder": img.get("subfolder", ""),
                    "type": "output"
                })
                if r.status_code == 200:
                    output_path.parent.mkdir(parents=True, exist_ok=True)
                    
                    # Load image and convert white background to transparent
                    img_data = Image.open(io.BytesIO(r.content)).convert("RGBA")
                    datas = img_data.getdata()
                    
                    new_data = []
                    for item in datas:
                        # If pixel is white (or near white), make transparent
                        if item[0] > 240 and item[1] > 240 and item[2] > 240:
                            new_data.append((255, 255, 255, 0))  # Transparent
                        else:
                            new_data.append(item)
                    
                    img_data.putdata(new_data)
                    img_data.save(str(output_path), "PNG")
                    return True
        return False
    except Exception as e:
        print(f"Download error: {e}")
        return False


def main():
    print("=" * 70)
    print("🎮 DOLINA SMRTI - ComfyUI Local Generator")
    print("   2D Indie Cartoon Style | Transparent Background")
    print("=" * 70)
    
    # Check ComfyUI
    try:
        r = requests.get(f"{COMFYUI_URL}/system_stats", timeout=5)
        v = r.json()["system"]["comfyui_version"]
        print(f"✅ ComfyUI v{v} running on port 8000")
    except:
        print("❌ ComfyUI not running!")
        print("   Start with: open /Applications/ComfyUI.app")
        return
    
    # Build registry
    print("\n📋 Building asset registry...")
    ASSETS = generate_registry()
    print(f"📊 Total: {len(ASSETS)} assets")
    
    cats = {}
    for a in ASSETS:
        cats[a["cat"]] = cats.get(a["cat"], 0) + 1
    print("\nBreakdown:")
    for cat, count in sorted(cats.items(), key=lambda x: -x[1]):
        print(f"  {cat}: {count}")
    
    existing = sum(1 for a in ASSETS if (OUTPUT_DIR / a["cat"] / a["file"]).exists())
    print(f"\n✅ Exist: {existing} | 🎯 Generate: {len(ASSETS) - existing}")
    
    print("\n🎮 STARTING GENERATION...\n")
    
    success, skip, fail = 0, 0, 0
    
    for i, asset in enumerate(ASSETS, 1):
        path = OUTPUT_DIR / asset["cat"] / asset["file"]
        
        if path.exists():
            skip += 1
            continue
        
        print(f"[{i}/{len(ASSETS)}] 🎨 {asset['file'][:40]}...")
        
        # Determine size
        size = 512
        if "tile_" in asset["file"]:
            size = 256  # Tiles are smaller
        elif "boss_" in asset["file"]:
            size = 768  # Bosses are larger
        
        wf = create_workflow(asset["prompt"], asset["file"].replace(".png", ""), size)
        pid = queue_prompt(wf)
        
        if pid and wait_completion(pid) and download_image(pid, path):
            print(f"         ✅ SAVED (transparent)")
            success += 1
        else:
            print(f"         ❌ FAILED")
            fail += 1
        
        if i % 10 == 0:
            print(f"\n📊 [{i}/{len(ASSETS)}] ✅{success} ⏭️{skip} ❌{fail}\n")
        
        time.sleep(0.5)
    
    print("\n" + "=" * 70)
    print("🎮 GENERATION COMPLETE!")
    print(f"✅ Success: {success} | ⏭️ Skip: {skip} | ❌ Fail: {fail}")
    print(f"📁 Output: {OUTPUT_DIR}")
    print("\nAll images have TRANSPARENT backgrounds (alpha channel)")


if __name__ == "__main__":
    main()