#!/usr/bin/env python3
"""
MASS ASSET GENERATOR - DolinaSmrti
Generates all 14,000+ game assets automatically using Vertex AI Imagen.
Follows the Dual-Style System and Production Workflow.
"""

import os
import sys
import json
import time
from pathlib import Path
from typing import Dict, List, Tuple
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/.asset_progress.json")

# Style definitions
STYLE_30 = """Style 30: Garden Story cozy adventure. Soft rounded friendly shapes, pastel-vibrant balance, wholesome mature aesthetic, Link's Awakening meets modern indie. 2D game asset. Solid chroma green background (#00FF00)."""

STYLE_32 = """Style 32: Cult of the Lamb cute-dark aesthetic. Bold outlines, noir shadows, pastel-gothic colors (dusty pinks, muted purples, cream whites), chunky proportions, big expressive eyes, contradiction between cute and dark. 2D game sprite. Solid chroma green background (#00FF00)."""

# Asset database - complete production plan
ASSET_DATABASE = {
    # PHASE 1: Plants (Style 30) - ALREADY COMPLETE
    "rastline": {
        "style": STYLE_30,
        "categories": {
            "sadje": [
                # Already generated - 30 assets
            ],
            "zelenjava": [
                # Already generated - 30 assets
            ],
            "ganja": [
                # Already generated - 15 assets
            ],
            "drevesa": [
                # Already generated - 10 assets
            ],
            "roze": [
                # Already generated - 8 assets
            ]
        }
    },
    
    # PHASE 2: Creatures (Style 32)
    "kreature": {
        "style": STYLE_32,
        "categories": {
            "zombiji": [
                ("zombie_basic", "Basic shambling zombie, rotting flesh"),
                ("zombie_runner", "Fast running zombie, aggressive pose"),
                ("zombie_tank", "Large bulky zombie, muscular"),
                ("zombie_spitter", "Zombie with acidic spit attack"),
                ("zombie_exploder", "Bloated zombie ready to explode"),
                ("zombie_crawler", "Crawling zombie on ground"),
                ("zombie_screamer", "Zombie with wide screaming mouth"),
                ("zombie_bride", "Zombie in tattered wedding dress"),
                ("zombie_cop", "Zombie police officer"),
                ("zombie_doctor", "Zombie in medical scrubs"),
            ],
            "dinozavri": [
                ("dino_trex", "T-Rex dinosaur, fierce predator"),
                ("dino_raptor", "Velociraptor, hunting pose"),
                ("dino_trike", "Triceratops, herbivore with horns"),
                ("dino_stego", "Stegosaurus with back plates"),
                ("dino_bronto", "Brontosaurus, long neck"),
                ("dino_ptero", "Pterodactyl flying"),
                ("dino_ankylo", "Ankylosaurus with club tail"),
                ("dino_parasaur", "Parasaurolophus with head crest"),
                ("dino_carno", "Carnotaurus with horns"),
                ("dino_spino", "Spinosaurus with sail"),
            ],
            "zivali": [
                ("animal_bear", "Brown bear, standing pose"),
                ("animal_wolf", "Gray wolf, hunting stance"),
                ("animal_deer", "Deer with antlers"),
                ("animal_rabbit", "Cute rabbit"),
                ("animal_fox", "Red fox, cunning"),
                ("animal_boar", "Wild boar, aggressive"),
                ("animal_crow", "Black crow perched"),
                ("animal_owl", "Wise owl"),
                ("animal_snake", "Coiled snake"),
                ("animal_rat", "Large rat"),
            ],
            "magicna_bitja": [
                ("magic_demon", "Small demon creature"),
                ("magic_imp", "Mischievous imp"),
                ("magic_wraith", "Dark wraith spirit"),
                ("magic_familiar", "Black cat familiar"),
                ("magic_golem", "Stone golem"),
            ]
        }
    },
    
    # PHASE 2: NPCs (Style 32)
    "npc": {
        "style": STYLE_32,
        "categories": {
            "policija": [
                ("police_chief", "Police chief, stern expression"),
                ("police_officer", "Regular police officer"),
                ("police_corrupt", "Corrupt cop, suspicious"),
            ],
            "zdravstvo": [
                ("doctor_main", "Town doctor, friendly"),
                ("nurse_helper", "Nurse assistant"),
            ],
            "trgovci": [
                ("merchant_general", "General store merchant"),
                ("merchant_seeds", "Seed merchant"),
                ("merchant_black_market", "Shady black market dealer"),
            ],
            "mentorji": [
                ("mentor_farming", "Old farmer mentor"),
                ("mentor_magic", "Dark magic mentor"),
            ]
        }
    },
    
    # PHASE 3: Main Characters (Style 32) - SAVE FOR LAST
    "glavni_karakterji": {
        "style": STYLE_32,
        "categories": {
            "kai": [
                ("kai_front_idle", "Kai facing front, idle stance"),
                ("kai_back_idle", "Kai facing back, idle stance"),
                ("kai_left_idle", "Kai facing left, idle stance"),
                ("kai_right_idle", "Kai facing right, idle stance"),
                ("kai_front_walk_1", "Kai front walk frame 1"),
                ("kai_front_walk_2", "Kai front walk frame 2"),
                ("kai_back_walk_1", "Kai back walk frame 1"),
                ("kai_back_walk_2", "Kai back walk frame 2"),
                ("kai_left_walk_1", "Kai left walk frame 1"),
                ("kai_left_walk_2", "Kai left walk frame 2"),
                ("kai_right_walk_1", "Kai right walk frame 1"),
                ("kai_right_walk_2", "Kai right walk frame 2"),
            ],
            "ana": [
                ("ana_front_idle", "Ana facing front, idle stance"),
                ("ana_back_idle", "Ana facing back, idle stance"),
                ("ana_left_idle", "Ana facing left, idle stance"),
                ("ana_right_idle", "Ana facing right, idle stance"),
                ("ana_front_walk_1", "Ana front walk frame 1"),
                ("ana_front_walk_2", "Ana front walk frame 2"),
                ("ana_back_walk_1", "Ana back walk frame 1"),
                ("ana_back_walk_2", "Ana back walk frame 2"),
                ("ana_left_walk_1", "Ana left walk frame 1"),
                ("ana_left_walk_2", "Ana left walk frame 2"),
                ("ana_right_walk_1", "Ana right walk frame 1"),
                ("ana_right_walk_2", "Ana right walk frame 2"),
            ],
            "gronk": [
                ("gronk_front_idle", "Gronk facing front, idle stance"),
                ("gronk_back_idle", "Gronk facing back, idle stance"),
                ("gronk_left_idle", "Gronk facing left, idle stance"),
                ("gronk_right_idle", "Gronk facing right, idle stance"),
            ],
            "susi": [
                ("susi_front_idle", "Susi the cat facing front"),
                ("susi_back_idle", "Susi the cat facing back"),
                ("susi_left_idle", "Susi the cat facing left"),
                ("susi_right_idle", "Susi the cat facing right"),
            ]
        }
    }
}


class MassAssetGenerator:
    """Generates all game assets systematically."""
    
    def __init__(self):
        """Initialize the generator."""
        print("🚀 DolinaSmrti Mass Asset Generator")
        print("=" * 60)
        
        # 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.total_generated = 0
        self.total_failed = 0
        self.session_generated = 0
        
    def load_progress(self) -> Dict:
        """Load generation progress from file."""
        if PROGRESS_FILE.exists():
            with open(PROGRESS_FILE, 'r') as f:
                return json.load(f)
        return {"completed": [], "failed": []}
    
    def save_progress(self):
        """Save current progress."""
        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 was already generated."""
        return asset_id in self.progress["completed"]
    
    def generate_image(self, prompt: str, filename: str, output_path: Path) -> bool:
        """Generate a single image using Vertex AI."""
        try:
            print(f"  Generating: {filename}...")
            
            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 image
                output_path.parent.mkdir(parents=True, exist_ok=True)
                response.images[0].save(location=str(output_path))
                print(f"    ✅ Saved: {output_path.name}")
                return True
            else:
                print(f"    ❌ No image generated")
                return False
                
        except Exception as e:
            print(f"    ❌ Error: {str(e)}")
            return False
    
    def generate_category(self, main_folder: str, category: str, 
                         assets: List[Tuple[str, str]], style: str):
        """Generate all assets in a category."""
        print(f"\n📁 {main_folder}/{category}")
        print("-" * 60)
        
        category_path = BASE_PATH / main_folder / category
        
        for filename, description in assets:
            asset_id = f"{main_folder}/{category}/{filename}"
            
            # Skip if already completed
            if self.is_completed(asset_id):
                print(f"  ⏭️  Skipping: {filename} (already generated)")
                continue
            
            # Build full prompt
            full_prompt = f"{style} {description}"
            
            # Generate
            output_file = category_path / f"{filename}.png"
            success = self.generate_image(full_prompt, filename, output_file)
            
            if success:
                self.progress["completed"].append(asset_id)
                self.total_generated += 1
                self.session_generated += 1
            else:
                self.progress["failed"].append(asset_id)
                self.total_failed += 1
            
            # Save progress after each asset
            self.save_progress()
            
            # Rate limiting - wait between generations
            time.sleep(1)
    
    def generate_all(self):
        """Generate all assets following the production workflow."""
        print("\n🎯 Starting mass generation...")
        print(f"Progress file: {PROGRESS_FILE}")
        print()
        
        start_time = time.time()
        
        # Follow production workflow order:
        # 1. Plants (already done, skip)
        # 2. Creatures & NPCs
        # 3. Main characters (last)
        
        # Phase 2: Creatures
        if "kreature" in ASSET_DATABASE:
            creatures = ASSET_DATABASE["kreature"]
            for category, assets in creatures["categories"].items():
                if assets:  # Only if there are assets defined
                    self.generate_category("kreature", category, assets, creatures["style"])
        
        # Phase 2: NPCs
        if "npc" in ASSET_DATABASE:
            npcs = ASSET_DATABASE["npc"]
            for category, assets in npcs["categories"].items():
                if assets:
                    self.generate_category("npc", category, assets, npcs["style"])
        
        # Phase 3: Main Characters (generate last for highest quality)
        if "glavni_karakterji" in ASSET_DATABASE:
            print("\n" + "=" * 60)
            print("⚠️  FINAL PHASE: Main Characters (Highest Quality)")
            print("=" * 60)
            chars = ASSET_DATABASE["glavni_karakterji"]
            for category, assets in chars["categories"].items():
                if assets:
                    self.generate_category("glavni_karakterji", category, assets, chars["style"])
        
        # Summary
        elapsed = time.time() - start_time
        print("\n" + "=" * 60)
        print("🎉 GENERATION COMPLETE!")
        print("=" * 60)
        print(f"Session generated: {self.session_generated}")
        print(f"Total generated: {self.total_generated}")
        print(f"Total failed: {self.total_failed}")
        print(f"Time elapsed: {elapsed/60:.1f} minutes")
        print(f"Progress saved to: {PROGRESS_FILE}")


def main():
    """Main entry point."""
    try:
        generator = MassAssetGenerator()
        generator.generate_all()
    except KeyboardInterrupt:
        print("\n\n⚠️  Generation interrupted by user")
        print("Progress has been saved. Run again to resume.")
        sys.exit(0)
    except Exception as e:
        print(f"\n❌ Fatal error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()