#!/usr/bin/env python3
"""
🎮 DOLINA SMRTI - ComfyUI Local Generator
Generates remaining assets with ComfyUI + rembg for transparent backgrounds
No API rate limits!
"""

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

# Install rembg on first run if needed
try:
    from rembg import remove
    from PIL import Image
except ImportError:
    print("Installing dependencies...")
    os.system("pip3 install --user rembg pillow onnxruntime")
    from rembg import remove
    from PIL import Image

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

# Style - matches the 2D Indie Cartoon style
STYLE = "2D indie game sprite, cartoon vector style, clean smooth lines, full body visible from head to feet standing on ground, tight crop, isolated on pure white background #FFFFFF"
NEGATIVE = "pixel art, pixelated, voxel, 3d, realistic, photo, anime, manga, chibi, blurry, watermark, text, green background"

# ============================================================================
# REMAINING ASSETS TO GENERATE
# ============================================================================

ASSETS = [
    # Remaining NPCs
    ("npcs", "npc_elder", "wise tribal elder NPC, walking stick, long white beard, robes"),
    ("npcs", "npc_child", "survivor orphan child NPC, oversized clothes, teddy bear"),
    ("npcs", "npc_cook", "camp cook NPC, pot and ladle, apron, big belly, friendly"),
    ("npcs", "npc_scout", "nomad scout NPC, binoculars, light leather armor, agile"),
    
    # Animation frames
    ("npcs", "gronk_front_walk1", "Gronk troll, PINK dreadlocks, black baggy shirt pants, front view, left leg forward walking"),
    ("npcs", "gronk_front_walk2", "Gronk troll, PINK dreadlocks, black baggy shirt pants, front view, right leg forward walking"),
    ("npcs", "gronk_back_idle", "Gronk troll, PINK dreadlocks, black baggy shirt pants, back view, standing idle"),
    ("npcs", "kai_front_walk1", "Kai survivor, GREEN dreadlocks, blue jacket, front view, left leg forward walking"),
    ("npcs", "kai_front_walk2", "Kai survivor, GREEN dreadlocks, blue jacket, front view, right leg forward walking"),
    ("npcs", "kai_back_idle", "Kai survivor, GREEN dreadlocks, blue jacket, back view, standing idle"),
    ("npcs", "ana_front_walk1", "Ana explorer, PINK dreadlocks, brown vest, front view, left leg forward walking"),
    ("npcs", "ana_front_walk2", "Ana explorer, PINK dreadlocks, brown vest, front view, right leg forward walking"),
    
    # Remaining animals
    ("zivali", "goat_brown", "brown goat farm animal, small horns, playful"),
    ("zivali", "rabbit_white", "white rabbit, long ears, fluffy"),
    ("zivali", "fox_red", "red fox wildlife, bushy tail, clever"),
    ("zivali", "golden_goose", "Golden Goose legendary, shimmering gold feathers"),
    
    # Remaining mutants
    ("mutanti", "slime_red", "red fire slime monster, hot flames"),
    ("mutanti", "slime_purple", "purple poison slime, toxic bubbling"),
    ("mutanti", "mutant_rat_giant", "giant rat mutant, dog sized, diseased"),
    ("mutanti", "zombie_bloated", "bloated zombie, huge swollen, toxic drool"),
    ("mutanti", "zombie_armored", "armored zombie, riot gear, tough"),
    ("mutanti", "zombie_crawler", "crawler zombie, no legs, crawling"),
    
    # Remaining bosses
    ("bosses", "boss_dragon_ice", "Ice Dragon boss, blue scales, frost breath"),
    ("bosses", "boss_phoenix", "Phoenix boss, giant fire bird, rebirth flames"),
    ("bosses", "boss_hydra", "Hydra boss, multi-headed serpent"),
    ("bosses", "boss_treant", "Elder Treant boss, giant living tree, nature"),
    ("bosses", "boss_demon", "Demon Prince boss, hellfire, horns"),
    
    # Remaining items
    ("items", "tool_pickaxe_iron", "iron pickaxe mining tool"),
    ("items", "tool_shovel", "shovel digging tool"),
    ("items", "tool_hammer", "hammer building tool"),
    ("items", "weapon_bow_wood", "wooden bow with arrows"),
    ("items", "weapon_spear", "wooden spear, hunting weapon"),
    ("items", "weapon_crossbow", "crossbow ranged weapon"),
    ("items", "wand_lightning", "lightning magic wand, yellow crystal, sparks"),
    ("items", "wand_earth", "earth magic wand, green crystal, nature"),
    ("items", "wand_shadow", "shadow magic wand, purple crystal, darkness"),
    ("items", "wand_healing", "healing wand, white crystal, glow"),
    ("items", "potion_mana", "mana potion, blue liquid bottle"),
    ("items", "backpack_small", "small leather backpack"),
    ("items", "helmet_metal", "metal protective helmet"),
    ("items", "gas_mask", "gas mask with filters"),
    ("items", "food_bread", "bread loaf, fresh baked"),
    ("items", "food_apple", "red apple, fresh fruit"),
    ("items", "food_meat", "cooked meat steak"),
    
    # Remaining environment
    ("environment", "tree_pine", "pine tree, evergreen, cone shaped"),
    ("environment", "tree_dead", "dead tree, no leaves, gnarled branches"),
    ("environment", "tree_cherry", "cherry blossom tree, pink flowers"),
    ("environment", "flower_red", "red flowers patch"),
    ("environment", "flower_blue", "blue flowers patch"),
    ("environment", "mushroom_red", "red spotted mushrooms"),
    ("environment", "fence_wood", "wooden fence section"),
    ("environment", "barrel_wood", "wooden barrel storage"),
    ("environment", "crate_wooden", "wooden shipping crate"),
    ("environment", "ruin_wall", "ruined brick wall section"),
    ("environment", "car_rusted", "rusted abandoned car wreck"),
    ("environment", "sign_warning", "warning sign post"),
]


def log(msg):
    ts = datetime.now().strftime("%H:%M:%S")
    print(f"[{ts}] {msg}")
    sys.stdout.flush()


def create_workflow(prompt_text, output_name, size=512):
    """Create ComfyUI workflow"""
    seed = int(time.time() * 1000) % 2147483647
    full_prompt = f"{STYLE}, {prompt_text}"
    
    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, "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):
    try:
        r = requests.post(
            f"{COMFYUI_URL}/prompt",
            json={"prompt": workflow, "client_id": f"ds_{uuid.uuid4().hex[:8]}"},
            timeout=10
        )
        return r.json().get("prompt_id")
    except Exception as e:
        log(f"❌ Queue error: {e}")
        return None


def wait_completion(prompt_id, timeout=120):
    start = time.time()
    while time.time() - start < timeout:
        try:
            r = requests.get(f"{COMFYUI_URL}/history/{prompt_id}", timeout=5)
            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_and_process(prompt_id, output_path):
    """Download from ComfyUI and remove background with rembg"""
    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 and remove background
                    from io import BytesIO
                    src_img = Image.open(BytesIO(r.content))
                    transparent_img = remove(src_img)
                    transparent_img.save(str(output_path), "PNG")
                    return True
        return False
    except Exception as e:
        log(f"❌ Download error: {e}")
        return False


def main():
    log("=" * 60)
    log("🎮 DOLINA SMRTI - ComfyUI Local Generator")
    log("   No API rate limits! All local processing!")
    log("=" * 60)
    
    # Check ComfyUI
    try:
        r = requests.get(f"{COMFYUI_URL}/system_stats", timeout=5)
        v = r.json()["system"]["comfyui_version"]
        log(f"✅ ComfyUI v{v} running")
    except:
        log("❌ ComfyUI not running! Start it first.")
        return
    
    # Get existing files
    existing = set()
    for cat, name, _ in ASSETS:
        path = OUTPUT_DIR / cat / f"{name}.png"
        if path.exists():
            existing.add(name)
    
    remaining = [(c, n, p) for c, n, p in ASSETS if n not in existing]
    
    log(f"\n📊 Status:")
    log(f"   Total in registry: {len(ASSETS)}")
    log(f"   Already generated: {len(existing)}")
    log(f"   Remaining: {len(remaining)}")
    
    if not remaining:
        log("\n✅ All assets already generated!")
        return
    
    log(f"\n🚀 Starting generation of {len(remaining)} assets...")
    log("   Each takes ~15-30 seconds\n")
    
    success, fail = 0, 0
    start_time = time.time()
    
    for i, (cat, name, prompt) in enumerate(remaining, 1):
        path = OUTPUT_DIR / cat / f"{name}.png"
        
        log(f"[{i}/{len(remaining)}] 🎨 {name}...")
        
        # Determine size
        size = 512
        if "boss_" in name:
            size = 768
        
        wf = create_workflow(prompt, name, size)
        pid = queue_prompt(wf)
        
        if pid and wait_completion(pid) and download_and_process(pid, path):
            log(f"         ✅ DONE (transparent)")
            success += 1
        else:
            log(f"         ❌ FAILED")
            fail += 1
        
        # Progress every 5
        if i % 5 == 0:
            elapsed = time.time() - start_time
            per_asset = elapsed / i
            remaining_time = per_asset * (len(remaining) - i)
            log(f"\n📊 Progress: {i}/{len(remaining)} | ✅{success} ❌{fail}")
            log(f"   ETA: {remaining_time/60:.1f} min\n")
    
    elapsed = time.time() - start_time
    log("\n" + "=" * 60)
    log("🎮 GENERATION COMPLETE!")
    log(f"   ✅ Success: {success}")
    log(f"   ❌ Failed: {fail}")
    log(f"   ⏱️  Time: {elapsed/60:.1f} min")
    log(f"   📁 Output: {OUTPUT_DIR}")
    log("=" * 60)


if __name__ == "__main__":
    main()