#!/usr/bin/env python3
"""
DOLINASMRTI PRODUCTION ASSET GENERATOR V2
Generates ALL game assets with proper Style 32/30 consistency.
Uses detailed prompts following the Dual-Style System Guide.
"""

import os
import json
import time
from pathlib import Path
from typing import Dict, List
import vertexai
from vertexai.preview.vision_models import ImageGenerationModel

# === CONFIGURATION ===
PROJECT_ID = "gen-lang-client-0428644398"
LOCATION = "us-central1"
BASE_PATH = Path("/Users/davidkotnik/repos/novafarma/assets/slike")
PROGRESS_FILE = Path("/Users/davidkotnik/repos/novafarma/.generation_progress.json")
LOG_FILE = Path("/Users/davidkotnik/repos/novafarma/generation.log")

# === STYLE DEFINITIONS ===

STYLE_32_BASE = """Style 32: Cult of the Lamb cute-dark aesthetic. Bold black outlines, pastel-gothic color palette (dusty pinks, muted purples, cream whites, soft grays with noir shadows), chunky proportions, big expressive eyes, smooth gradients, contradiction between adorable and sinister. """

STYLE_30_BASE = """Style 30: Garden Story cozy adventure. Soft rounded friendly shapes, pastel-vibrant color balance (cheerful but not childish), wholesome mature aesthetic, Link's Awakening meets modern indie charm, clean vector look with smooth edges. """

COMMON_SUFFIX = """2D game asset, centered subject with 10px margin, consistent lighting. Solid chroma green background (#00FF00) for easy removal."""

# === DETAILED PROMPT DATABASE ===

PHASE_2_CREATURES = {
    "kreature/zombiji": {
        "style": "style_32",
        "assets": [
            ("zombie_basic", f"{STYLE_32_BASE}Basic shambling zombie, rotting flesh with exposed ribs, tattered clothing, blank white eyes, slow stumbling pose, covered in dirt and blood splatters, cute-grotesque balance. {COMMON_SUFFIX}"),
            ("zombie_runner", f"{STYLE_32_BASE}Fast running zombie, aggressive sprinting pose, lean athletic build, wild disheveled hair, crazed expression, torn sportswear, speed lines effect. {COMMON_SUFFIX}"),
            ("zombie_tank", f"{STYLE_32_BASE}Large bulky zombie, massive muscular frame, cracked skin showing bones, heavy stomping stance, intimidating but cute chunky proportions, tattered vest. {COMMON_SUFFIX}"),
            ("zombie_spitter", f"{STYLE_32_BASE}Zombie with acidic spit attack, bloated throat pouch, green acid dripping from mouth, hunched pose ready to spit, pustules on skin, toxic cute. {COMMON_SUFFIX}"),
            ("zombie_exploder", f"{STYLE_32_BASE}Bloated zombie ready to explode, swollen round belly with glowing veins, worried expression, arms spread wide, pulsing red glow inside body, ticking time bomb cuteness. {COMMON_SUFFIX}"),
            ("zombie_crawler", f"{STYLE_32_BASE}Crawling zombie on ground, dragging itself forward with arms, legless lower body, determined expression, dirty fingernails, trail of slime. {COMMON_SUFFIX}"),
            ("zombie_screamer", f"{STYLE_32_BASE}Screaming zombie with wide open mouth, sound wave visual effect, hands clutching head, red glowing throat, ear-splitting pose, disturbing yet adorable. {COMMON_SUFFIX}"),
            ("zombie_bride", f"{STYLE_32_BASE}Zombie in tattered wedding dress, decayed bride with veil, bouquet of dead flowers, melancholy expression, torn white gown with blood stains, tragic beauty. {COMMON_SUFFIX}"),
            ("zombie_cop", f"{STYLE_32_BASE}Zombie police officer, torn blue uniform with badge, police cap askew, holding broken baton, deadpan law enforcement expression, duty-bound undead. {COMMON_SUFFIX}"),
            ("zombie_doctor", f"{STYLE_32_BASE}Zombie in medical scrubs, stethoscope around neck, blood-stained white coat, surgical mask hanging loose, carrying clipboard, twisted healthcare worker. {COMMON_SUFFIX}"),
        ]
    },
    
    "kreature/dinozavri": {
        "style": "style_32",
        "assets": [
            ("dino_trex", f"{STYLE_32_BASE}T-Rex dinosaur, fierce predator with tiny arms, massive head with sharp teeth showing, powerful legs in hunting stance, scaly texture, roaring pose, chunky cute-terrifying proportions. {COMMON_SUFFIX}"),
            ("dino_raptor", f"{STYLE_32_BASE}Velociraptor, intelligent hunting pose, sickle claw prominently raised, sleek body with feathered accents, cunning big eyes, pack hunter stance. {COMMON_SUFFIX}"),
            ("dino_trike", f"{STYLE_32_BASE}Triceratops, herbivore with three prominent horns and neck frill, sturdy body, gentle but powerful stance, armored head, grazing pose. {COMMON_SUFFIX}"),
            ("dino_stego", f"{STYLE_32_BASE}Stegosaurus with distinctive back plates in row, spiked tail (thagomizer), long neck extended, plant-eater pose, armored gentle giant. {COMMON_SUFFIX}"),
            ("dino_bronto", f"{STYLE_32_BASE}Brontosaurus, extremely long neck reaching upward, massive body on four thick legs, small head, gentle giant expression, eating tree leaves. {COMMON_SUFFIX}"),
            ("dino_ankylo", f"{STYLE_32_BASE}Ankylosaurus with heavy armor plating covering body, club tail weapon raised defensively, low waddle stance, tank-like build, prehistoric tank. {COMMON_SUFFIX}"),
            ("dino_ptero", f"{STYLE_32_BASE}Pterodactyl flying dinosaur, wings spread wide in flight, long beak with sharp teeth, soaring pose, leather wings, sky hunter. {COMMON_SUFFIX}"),
            ("dino_spino", f"{STYLE_32_BASE}Spinosaurus with massive sail on back, crocodile-like snout, semi-aquatic stance, powerful arms, fishing predator pose. {COMMON_SUFFIX}"),
            ("dino_dilo", f"{STYLE_32_BASE}Dilophosaurus with distinctive twin head crests, neck frill extended in threat display, venom-spitting pose, medium-sized predator. {COMMON_SUFFIX}"),
            ("dino_pachy", f"{STYLE_32_BASE}Pachycephalosaurus with thick domed skull for head-butting, bipedal stance, charging pose, bony head visible, ram attack ready. {COMMON_SUFFIX}"),
        ]
    },
    
    "kreature/zivali": {
        "style": "style_32",
        "assets": [
            ("animal_bear", f"{STYLE_32_BASE}Brown bear standing on hind legs, powerful build, thick fur, cute rounded ears, big paws with claws, intimidating but huggable stance. {COMMON_SUFFIX}"),
            ("animal_wolf", f"{STYLE_32_BASE}Gray wolf in hunting stance, sleek body lowered ready to pounce, sharp teeth showing, pointed ears alert, pack leader aura, fierce yet cute. {COMMON_SUFFIX}"),
            ("animal_deer", f"{STYLE_32_BASE}Deer with majestic antlers, gentle eyes, standing gracefully, spotted coat, alert ears, forest spirit vibe, serene but watchful. {COMMON_SUFFIX}"),
            ("animal_rabbit", f"{STYLE_32_BASE}Cute rabbit sitting upright, long ears perked up, fluffy tail, whiskers twitching, holding carrot, innocent big eyes, forest cottontail. {COMMON_SUFFIX}"),
            ("animal_fox", f"{STYLE_32_BASE}Red fox with bushy tail, cunning expression, sleek orange fur with white chest, pointed snout, sly sitting pose, forest trickster. {COMMON_SUFFIX}"),
            ("animal_boar", f"{STYLE_32_BASE}Wild boar, aggressive stance with head lowered, sharp tusks protruding, bristly fur, muscular build, charging pose, dangerous but cute. {COMMON_SUFFIX}"),
            ("animal_crow", f"{STYLE_32_BASE}Black crow perched on branch, glossy black feathers, intelligent beady eyes, sharp beak, one foot raised, ominous messenger. {COMMON_SUFFIX}"),
            ("animal_owl", f"{STYLE_32_BASE}Wise owl perched, large round eyes staring forward, feathered ear tufts, soft brown plumage, nocturnal hunter pose, mysterious knowledge-keeper. {COMMON_SUFFIX}"),
            ("animal_snake", f"{STYLE_32_BASE}Coiled snake ready to strike, forked tongue flicking out, scaled body in spiral, diamond head pattern, venomous but cute, danger noodle. {COMMON_SUFFIX}"),
            ("animal_rat", f"{STYLE_32_BASE}Large rat on hind legs, long naked tail curled, whiskers prominent, beady eyes, tiny paws, city scavenger pose, plague cute. {COMMON_SUFFIX}"),
        ]
    },
    
    "kreature/magicna_bitja": {
        "style": "style_32",
        "assets": [
            ("magic_demon", f"{STYLE_32_BASE}Small demon creature, red skin with tiny horns, pointed tail, pitchfork weapon, mischievous grin, clawed feet, adorably evil. {COMMON_SUFFIX}"),
            ("magic_imp", f"{STYLE_32_BASE}Mischievous imp with bat wings, green goblin skin, pointy ears, sharp teeth grin, crouching pose, troublemaker aura. {COMMON_SUFFIX}"),
            ("magic_wraith", f"{STYLE_32_BASE}Dark wraith spirit, shadowy hooded figure with no face visible, tattered robes flowing, skeletal hands extending, floating pose with wispy trails, death incarnate. {COMMON_SUFFIX}"),
            ("magic_familiar", f"{STYLE_32_BASE}Black cat familiar with glowing yellow eyes, sleek fur, mystical aura, sitting with tail curled, witch's companion, magical energy swirls. {COMMON_SUFFIX}"),
            ("magic_golem", f"{STYLE_32_BASE}Stone golem, heavy blocky rock body, runes glowing on chest, moss growing on shoulders, powerful fists, slow lumbering stance, ancient guardian. {COMMON_SUFFIX}"),
        ]
    },
    
    "npc/policija": {
        "style": "style_32",
        "assets": [
            ("police_chief", f"{STYLE_32_BASE}Police chief NPC, stern middle-aged person in decorated uniform, police cap with badge, medals on chest, hands on hips authoritative pose, graying hair, law enforcement leader. {COMMON_SUFFIX}"),
            ("police_officer", f"{STYLE_32_BASE}Regular police officer NPC, standard blue uniform, police cap, badge visible, utility belt with radio, friendly but professional expression, community protector. {COMMON_SUFFIX}"),
            ("police_corrupt", f"{STYLE_32_BASE}Corrupt cop NPC, disheveled uniform, suspicious shifty eyes, sinister grin, holding bribe money behind back, morally gray police officer. {COMMON_SUFFIX}"),
        ]
    },
    
    "npc/zdravstvo": {
        "style": "style_32",
        "assets": [
            ("doctor_main", f"{STYLE_32_BASE}Town doctor NPC, friendly physician in white coat, stethoscope around neck, kind wise expression, glasses, holding medical clipboard, trusted healer. {COMMON_SUFFIX}"),
            ("nurse_helper", f"{STYLE_32_BASE}Nurse assistant NPC, medical scrubs and cap, caring expression, holding first aid kit, supportive pose, healthcare worker. {COMMON_SUFFIX}"),
        ]
    },
    
    "npc/trgovci": {
        "style": "style_32",
        "assets": [
            ("merchant_general", f"{STYLE_32_BASE}General store merchant NPC, apron over casual clothes, welcoming smile, holding product, behind counter pose, friendly shopkeeper. {COMMON_SUFFIX}"),
            ("merchant_seeds", f"{STYLE_32_BASE}Seed merchant NPC, farmer's clothes with seed packets, straw hat, holding vegetable seeds, dirt on hands, agricultural expert. {COMMON_SUFFIX}"),
            ("merchant_black_market", f"{STYLE_32_BASE}Shady black market dealer NPC, dark hood concealing face, mysterious aura, holding questionable goods, lurking in shadows, underground trader. {COMMON_SUFFIX}"),
        ]
    },
    
    "npc/mentorji": {
        "style": "style_32",
        "assets": [
            ("mentor_farming", f"{STYLE_32_BASE}Old farmer mentor NPC, elderly with weathered face, overalls and straw hat, wise smile, leaning on hoe, years of experience visible. {COMMON_SUFFIX}"),
            ("mentor_magic", f"{STYLE_32_BASE}Dark magic mentor NPC, robed figure with mystical symbols, glowing hands with magical energy, ancient tome under arm, sinister wisdom. {COMMON_SUFFIX}"),
        ]
    },
}


class ProductionAssetGenerator:
    """Professional asset generator with proper style consistency."""
    
    def __init__(self):
        print("=" * 70)
        print("🎨 DOLINASMRTI PRODUCTION ASSET GENERATOR V2")
        print("=" * 70)
        print()
        
        # Initialize Vertex AI
        vertexai.init(project=PROJECT_ID, location=LOCATION)
        self.model = ImageGenerationModel.from_pretrained("imagen-3.0-generate-001")
        
        # Load progress
        self.progress = self.load_progress()
        
        # Stats
        self.session_count = 0
        self.session_failed = 0
        
        # Open log file
        self.log = open(LOG_FILE, 'a')
        self.log.write(f"\n\n{'='*70}\n")
        self.log.write(f"Session started: {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
        self.log.write(f"{'='*70}\n\n")
    
    def __del__(self):
        """Cleanup."""
        if hasattr(self, 'log'):
            self.log.close()
    
    def load_progress(self) -> Dict:
        """Load generation progress."""
        if PROGRESS_FILE.exists():
            with open(PROGRESS_FILE, 'r') as f:
                return json.load(f)
        return {"completed": [], "failed": [], "total_generated": 0}
    
    def save_progress(self):
        """Save current progress."""
        self.progress["total_generated"] = len(self.progress["completed"])
        with open(PROGRESS_FILE, 'w') as f:
            json.dump(self.progress, f, indent=2)
    
    def is_completed(self, asset_id: str) -> bool:
        """Check if asset already generated."""
        return asset_id in self.progress["completed"]
    
    def log_message(self, message: str):
        """Write to log file and console."""
        print(message)
        self.log.write(message + "\n")
        self.log.flush()
    
    def generate_asset(self, category_path: str, filename: str, prompt: str) -> bool:
        """Generate single asset with retry logic."""
        asset_id = f"{category_path}/{filename}"
        
        # Skip if completed
        if self.is_completed(asset_id):
            self.log_message(f"  ⏭️  SKIP: {filename} (already done)")
            return True
        
        self.log_message(f"  🎨 GENERATING: {filename}")
        self.log_message(f"     Prompt: {prompt[:100]}...")
        
        try:
            # Generate image
            response = self.model.generate_images(
                prompt=prompt,
                number_of_images=1,
                aspect_ratio="1:1",
                safety_filter_level="block_few",
                person_generation="allow_adult"
            )
            
            if response.images:
                # Save to correct folder
                output_dir = BASE_PATH / category_path
                output_dir.mkdir(parents=True, exist_ok=True)
                output_path = output_dir / f"{filename}.png"
                
                response.images[0].save(location=str(output_path))
                
                self.log_message(f"     ✅ SUCCESS: Saved to {output_path}")
                self.progress["completed"].append(asset_id)
                self.session_count += 1
                self.save_progress()
                return True
            else:
                self.log_message(f"     ❌ FAILED: No image returned")
                self.progress["failed"].append(asset_id)
                self.session_failed += 1
                self.save_progress()
                return False
                
        except Exception as e:
            self.log_message(f"     ❌ ERROR: {str(e)}")
            self.progress["failed"].append(asset_id)
            self.session_failed += 1
            self.save_progress()
            return False
    
    def generate_category(self, category_path: str, assets: List):
        """Generate all assets in a category."""
        self.log_message(f"\n{'='*70}")
        self.log_message(f"📁 CATEGORY: {category_path}")
        self.log_message(f"{'='*70}\n")
        
        for filename, prompt in assets:
            self.generate_asset(category_path, filename, prompt)
            
            # Rate limiting
            time.sleep(1.5)
    
    def generate_phase_2(self):
        """Generate Phase 2: Creatures & NPCs."""
        self.log_message("\n" + "🔥"*35)
        self.log_message("PHASE 2: CREATURES & NPCs (Style 32)")
        self.log_message("🔥"*35 + "\n")
        
        for category, data in PHASE_2_CREATURES.items():
            self.generate_category(category, data["assets"])
        
        self.log_message(f"\n{'='*70}")
        self.log_message("✅ PHASE 2 COMPLETE!")
        self.log_message(f"   Session generated: {self.session_count}")
        self.log_message(f"   Session failed: {self.session_failed}")
        self.log_message(f"   Total complete: {len(self.progress['completed'])}")
        self.log_message(f"{'='*70}\n")


def main():
    """Main entry point."""
    generator = ProductionAssetGenerator()
    
    try:
        generator.generate_phase_2()
    except KeyboardInterrupt:
        print("\n\n⚠️  Interrupted - Progress saved!")
    except Exception as e:
        print(f"\n❌ Fatal error: {e}")
        raise


if __name__ == "__main__":
    main()