#!/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
import subprocess
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 git_commit(file_path, category, name):
    """Auto-commit generated asset to Git"""
    try:
        repo_dir = Path("/Users/davidkotnik/repos/novafarma")
        rel_path = file_path.relative_to(repo_dir)
        
        # Git add
        subprocess.run(
            ["git", "add", str(rel_path)],
            cwd=repo_dir,
            check=True,
            capture_output=True
        )
        
        # Git commit
        commit_msg = f"🎨 Auto-generated asset: {category}/{name}"
        result = subprocess.run(
            ["git", "commit", "-m", commit_msg],
            cwd=repo_dir,
            capture_output=True,
            text=True
        )
        
        if result.returncode == 0:
            log(f"         📝 Git committed")
            return True
        else:
            # Možno da je že committed ali ni sprememb
            return False
    except Exception as e:
        log(f"         ⚠️  Git: {str(e)[:50]}")
        return False


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)")
            # Auto-commit to Git
            git_commit(path, cat, name)
            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()