#!/usr/bin/env python3
"""
DOLINASMRTI SPRITESHEET GENERATOR V3
Generates complete animation spritesheets for all game entities.
Style 32 for characters/creatures, Style 30 for plants.
Stardew Valley inspired animation system.
"""

import os
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/.spritesheet_progress.json")
LOG_FILE = Path("/Users/davidkotnik/repos/novafarma/spritesheet_generation.log")

# === STYLE 32: CULT OF THE LAMB (PROPER SPECS) ===
STYLE_32_DETAILED = """Style 32: Cult of the Lamb cute-dark aesthetic. CRITICAL REQUIREMENTS: Bold thick black outlines (2-3px), chibi proportions (head = 1/3 of body height), pastel-gothic color palette (dusty pink #E8B4D4, muted purple #9B7FB5, cream white #F5EFE6, soft gray #B8B8B8 with noir black shadows #2D2D2D), big expressive eyes (1/4 of face), smooth gradients on rounded forms, contradiction between adorable shapes and dark themes. Stardew Valley inspired pixel-perfect clarity. """

# === STYLE 30: GARDEN STORY COZY ===
STYLE_30_DETAILED = """Style 30: Garden Story cozy adventure. Soft rounded organic shapes, pastel-vibrant balance (cheerful greens #A8D5BA, warm yellows #F7DC6F, fresh oranges #F8B739), wholesome mature aesthetic (not childish), Link's Awakening charm with modern indie polish, clean vector edges with subtle texture. """

COMMON_SUFFIX = """2D game sprite, pixel-perfect clarity, centered composition with 10px margin, consistent top-down lighting, solid chroma green background (#00FF00)."""

# === SPRITESHEET DEFINITIONS ===

CREATURE_SPRITESHEETS = {
    "kreature/zombiji": {
        "style": STYLE_32_DETAILED,
        "entities": [
            {
                "name": "zombie_basic",
                "description": "Basic shambling zombie, rotting flesh, exposed ribs, tattered clothing, blank white eyes, slow gait",
                "animations": {
                    "idle": "Standing idle, slight body sway, 2 frames",
                    "walk_south": "Walking toward camera, shambling gait, 4 frames",
                    "walk_north": "Walking away, shambling, 4 frames", 
                    "walk_east": "Walking right, side view shamble, 4 frames",
                    "walk_west": "Walking left, mirror of east walk, 4 frames",
                    "attack": "Lunge forward with arms extended, bite motion, 3 frames"
                }
            },
            {
                "name": "zombie_runner",
                "description": "Fast aggressive zombie, sprinting pose, lean athletic build, wild hair, crazed expression",
                "animations": {
                    "idle": "Aggressive stance, heavy breathing, 2 frames",
                    "walk_south": "Running sprint toward camera, 4 frames",
                    "walk_north": "Sprinting away, 4 frames",
                    "walk_east": "Fast run right, 4 frames",
                    "walk_west": "Fast run left, 4 frames",
                    "attack": "Diving tackle motion, 3 frames"
                }
            },
            {
                "name": "zombie_tank",
                "description": "Massive bulky zombie, muscular frame, cracked skin, intimidating but cute chunky",
                "animations": {
                    "idle": "Heavy breathing, arms at sides, 2 frames",
                    "walk_south": "Slow powerful stomp forward, 4 frames",
                    "walk_north": "Lumbering away, 4 frames",
                    "walk_east": "Heavy walk right, 4 frames",
                    "walk_west": "Heavy walk left, 4 frames",
                    "attack": "Overhead smash motion, 3 frames"
                }
            }
        ]
    },
    
    "kreature/dinozavri": {
        "style": STYLE_32_DETAILED,
        "entities": [
            {
                "name": "dino_raptor",
                "description": "Velociraptor, intelligent predator, sickle claw raised, sleek feathered body, cunning eyes",
                "animations": {
                    "idle": "Alert stance, head bobbing, 2 frames",
                    "walk_south": "Hopping gait toward camera, Stardew Valley style, 4 frames",
                    "walk_north": "Hopping away, tail balance, 4 frames",
                    "walk_east": "Side-view hop right, 4 frames",
                    "walk_west": "Side-view hop left, 4 frames",
                    "attack": "Sickle claw slash motion, 3 frames"
                }
            },
            {
                "name": "dino_trex",
                "description": "T-Rex, massive head, tiny arms, powerful legs, roaring pose, chunky cute-terrifying",
                "animations": {
                    "idle": "Standing tall, slight head sway, 2 frames",
                    "walk_south": "Heavy stomp forward, ground shake, 4 frames",
                    "walk_north": "Stomping away, 4 frames",
                    "walk_east": "Side stomp right, 4 frames",
                    "walk_west": "Side stomp left, 4 frames",
                    "attack": "Massive bite lunge, jaws wide, 3 frames"
                }
            }
        ]
    },
    
    "kreature/zivali": {
        "style": STYLE_32_DETAILED,
        "entities": [
            {
                "name": "animal_bear",
                "description": "Brown bear, thick fur, rounded ears, big paws, intimidating but huggable",
                "animations": {
                    "idle": "Standing on hind legs, gentle sway, 2 frames",
                    "walk_south": "Four-legged lumbering walk forward, 4 frames",
                    "walk_north": "Lumbering away, 4 frames",
                    "walk_east": "Side walk right, 4 frames",
                    "walk_west": "Side walk left, 4 frames",
                    "attack": "Swipe with paw, standing tall, 3 frames"
                }
            },
            {
                "name": "animal_wolf",
                "description": "Gray wolf, sleek predator, sharp teeth, pointed ears, fierce yet cute",
                "animations": {
                    "idle": "Alert stance, ears perked, 2 frames",
                    "walk_south": "Prowling walk forward, 4 frames",
                    "walk_north": "Prowl away, 4 frames",
                    "walk_east": "Side prowl right, 4 frames",
                    "walk_west": "Side prowl left, 4 frames",
                    "attack": "Lunge bite motion, 3 frames"
                }
            }
        ]
    }
}


class SpritesheetGenerator:
    """Generates complete animation spritesheets for game entities."""
    
    def __init__(self):
        print("=" * 70)
        print("🎮 DOLINASMRTI SPRITESHEET GENERATOR V3")
        print("=" * 70)
        print("Features: Full animation support, Style 32 accurate, Auto-organization")
        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
        
        # Logging
        self.log = open(LOG_FILE, 'a')
        self.log_message(f"\n{'='*70}")
        self.log_message(f"Session started: {time.strftime('%Y-%m-%d %H:%M:%S')}")
        self.log_message(f"{'='*70}\n")
    
    def __del__(self):
        if hasattr(self, 'log'):
            self.log.close()
    
    def load_progress(self) -> Dict:
        if PROGRESS_FILE.exists():
            with open(PROGRESS_FILE, 'r') as f:
                return json.load(f)
        return {"completed": [], "failed": []}
    
    def save_progress(self):
        with open(PROGRESS_FILE, 'w') as f:
            json.dump(self.progress, f, indent=2)
    
    def is_completed(self, asset_id: str) -> bool:
        return asset_id in self.progress["completed"]
    
    def log_message(self, msg: str):
        print(msg)
        self.log.write(msg + "\n")
        self.log.flush()
    
    def generate_animation_frame(self, entity_name: str, animation_name: str, 
                                 description: str, style: str) -> bool:
        """Generate single animation frame."""
        
        asset_id = f"{entity_name}_{animation_name}"
        
        if self.is_completed(asset_id):
            self.log_message(f"    ⏭️  SKIP: {animation_name} (done)")
            return True
        
        # Build complete prompt
        full_prompt = f"{style}{description}, {animation_name}. {COMMON_SUFFIX}"
        
        self.log_message(f"    🎨 {animation_name}")
        
        try:
            response = self.model.generate_images(
                prompt=full_prompt,
                number_of_images=1,
                aspect_ratio="1:1",
                safety_filter_level="block_few",
                person_generation="allow_adult"
            )
            
            if response.images:
                # Save frame
                # Extract category from entity metadata (would need to pass this)
                # For now, save to temp location
                output_path = Path(f"/tmp/{entity_name}_{animation_name}.png")
                response.images[0].save(location=str(output_path))
                
                self.log_message(f"       ✅ Saved")
                self.progress["completed"].append(asset_id)
                self.session_count += 1
                self.save_progress()
                return True
            else:
                self.log_message(f"       ❌ No image")
                self.progress["failed"].append(asset_id)
                self.session_failed += 1
                return False
                
        except Exception as e:
            self.log_message(f"       ❌ ERROR: {str(e)}")
            self.progress["failed"].append(asset_id)
            self.session_failed += 1
            return False
    
    def generate_entity(self, category: str, entity: Dict, style: str):
        """Generate complete spritesheet for one entity."""
        
        name = entity["name"]
        desc = entity["description"]
        animations = entity["animations"]
        
        self.log_message(f"\n  🎯 ENTITY: {name}")
        self.log_message(f"     {desc}")
        self.log_message(f"     Animations: {len(animations)}")
        
        for anim_name, anim_desc in animations.items():
            full_desc = f"{desc}, {anim_desc}"
            self.generate_animation_frame(name, anim_name, full_desc, style)
            time.sleep(2)  # Rate limiting
    
    def generate_all_spritesheets(self):
        """Generate all entity spritesheets."""
        
        self.log_message("\n🔥 STARTING SPRITESHEET GENERATION")
        self.log_message("="*70 + "\n")
        
        for category, data in CREATURE_SPRITESHEETS.items():
            self.log_message(f"\n📁 CATEGORY: {category}")
            self.log_message("="*70)
            
            for entity in data["entities"]:
                self.generate_entity(category, entity, data["style"])
        
        # Summary
        self.log_message(f"\n{'='*70}")
        self.log_message("✅ GENERATION COMPLETE!")
        self.log_message(f"   Session generated: {self.session_count}")
        self.log_message(f"   Session failed: {self.session_failed}")
        self.log_message(f"{'='*70}\n")


def main():
    generator = SpritesheetGenerator()
    
    try:
        generator.generate_all_spritesheets()
    except KeyboardInterrupt:
        print("\n⚠️  Interrupted - Progress saved!")
    except Exception as e:
        print(f"\n❌ Error: {e}")
        raise


if __name__ == "__main__":
    main()