#!/usr/bin/env python3
"""
🌙 NOČNA GENERACIJA - FULL Asset Set
BREZ NPCs in glavnih karakterjev
Z avtomatskim Git commitom in progress trackingom
"""

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

try:
    from rembg import remove
    from PIL import Image
except ImportError:
    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")
REPO_DIR = Path("/Users/davidkotnik/repos/novafarma")
PROGRESS_LOG = REPO_DIR / "generation_progress.log"

STYLE = """Clean cartoon vector game sprite, bold smooth outlines, flat vibrant colors, 
minimal shading, cute playful aesthetic, full body visible head-to-toe, 
isolated on pure white background, 2D game art similar to Don't Starve or Zelda Link's Awakening style"""

NEGATIVE = """pixel art, pixelated, voxel, 3d render, realistic, photorealistic, 
detailed shading, gradient heavy, anime style, manga, blurry, low quality, 
watermark, text, green background, Stardew Valley style, farm sim generic style"""

def log(msg):
    ts = datetime.now().strftime("%H:%M:%S")
    line = f"[{ts}] {msg}"
    print(line)
    sys.stdout.flush()
    with open(PROGRESS_LOG, 'a') as f:
        f.write(line + "\n")

def git_commit(file_path, category, name):
    try:
        rel_path = file_path.relative_to(REPO_DIR)
        subprocess.run(["git", "add", str(rel_path)], cwd=REPO_DIR, check=True, capture_output=True)
        subprocess.run(["git", "commit", "-m", f"🎨 Generated: {category}/{name}"], cwd=REPO_DIR, capture_output=True)
        return True
    except:
        return False

# Load FULL registry safely
sys.path.insert(0, str(REPO_DIR / "scripts"))
import generate_assets_full
ALL_ASSETS = generate_assets_full.generate_registry()
FILTERED_ASSETS = [a for a in ALL_ASSETS if a['cat'] != 'npcs']

def create_workflow(prompt_text, output_name, size=512):
    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"full_{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)
            if prompt_id in r.json() and r.json()[prompt_id].get("outputs"):
                return True
        except:
            pass
        time.sleep(2)
    return False

def download_and_process(prompt_id, output_path):
    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)
                    from io import BytesIO
                    transparent_img = remove(Image.open(BytesIO(r.content)))
                    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("="*70)
    log("🌙 DAYTIME GENERATION - FULL Asset Set")
    log("   BREZ NPCs in glavnih karakterjev")
    log("="*70)
    
    try:
        r = requests.get(f"{COMFYUI_URL}/system_stats", timeout=5)
        log(f"✅ ComfyUI v{r.json()['system']['comfyui_version']}")
    except:
        log("❌ ComfyUI not running!")
        return
    
    log(f"\n📊 Registry Statistics:")
    log(f"   Total assets: {len(ALL_ASSETS)}")
    log(f"   NPCs (excluded): {len([a for a in ALL_ASSETS if a['cat'] == 'npcs'])}")
    log(f"   FILTERED (to generate): {len(FILTERED_ASSETS)}")
    
    existing = sum(1 for a in FILTERED_ASSETS if (OUTPUT_DIR / a['cat'] / a['file']).exists())
    to_generate = len(FILTERED_ASSETS) - existing
    
    log(f"\n📁 File Status:")
    log(f"   Already exist: {existing}")
    log(f"   To generate: {to_generate}")
    
    log(f"\n📦 By Category:")
    cats = {}
    for a in FILTERED_ASSETS:
        cats[a['cat']] = cats.get(a['cat'], 0) + 1
    for cat, count in sorted(cats.items(), key=lambda x: -x[1]):
        existing_cat = sum(1 for a in FILTERED_ASSETS if a['cat'] == cat and (OUTPUT_DIR / a['cat'] / a['file']).exists())
        log(f"   {cat:20s}: {count:4d} total ({existing_cat} exist, {count-existing_cat} to gen)")
    
    if to_generate == 0:
        log("\n✅ All assets already generated!")
        return
    
    log(f"\n🚀 Starting generation of {to_generate} assets...")
    log(f"   ETA: ~{to_generate * 20 / 3600:.1f} hours\n")
    
    success, skip, fail = 0, 0, 0
    start_time = time.time()
    
    for i, asset in enumerate(FILTERED_ASSETS, 1):
        path = OUTPUT_DIR / asset['cat'] / asset['file']
        
        if path.exists():
            skip += 1
            if i % 50 == 0:
                log(f"[{i}/{len(FILTERED_ASSETS)}] Progress: ✅{success} ⏭️ {skip} ❌{fail}")
            continue
        
        log(f"[{i}/{len(FILTERED_ASSETS)}] 🎨 {asset['cat']}/{asset['file'][:40]}")
        
        size = 768 if asset['cat'] == 'bosses' else 512
        wf = create_workflow(asset['prompt'], asset['file'].replace('.png', ''), size)
        pid = queue_prompt(wf)
        
        if pid and wait_completion(pid) and download_and_process(pid, path):
            log(f"             ✅ Done")
            git_commit(path, asset['cat'], asset['file'])
            success += 1
        else:
            log(f"             ❌ FAILED")
            fail += 1
        
        if i % 10 == 0:
            elapsed = time.time() - start_time
            per_asset = elapsed / (success + fail) if (success + fail) > 0 else 20
            remaining = per_asset * (to_generate - success - fail)
            log(f"\n📊 [{i}/{len(FILTERED_ASSETS)}] ✅{success} ⏭️ {skip} ❌{fail}")
            log(f"   ⏱️  Elapsed: {elapsed/60:.1f}m | ETA: {remaining/60:.1f}m\n")
    
    elapsed = time.time() - start_time
    log("\n" + "="*70)
    log("🎨 GENERATION COMPLETE!")
    log(f"   ✅ Success: {success}")
    log(f"   ⏭️  Skipped: {skip}")
    log(f"   ❌ Failed: {fail}")
    log(f"   ⏱️  Total time: {elapsed/3600:.1f} hours")
    log(f"   📁 Output: {OUTPUT_DIR}")
    log("="*70)

if __name__ == "__main__":
    main()