#!/usr/bin/env python3
"""
Master Asset Processor - Mrtva Dolina
Processes all game assets:
- Remove green screen (#00FF00) → transparent
- Separate multi-object images
- Organize into clean structure
"""

import os
import json
from pathlib import Path
from PIL import Image
import numpy as np
from scipy import ndimage
import shutil

class MrtvaDolina_AssetProcessor:
    def __init__(self, base_dir="/Users/davidkotnik/Desktop/novafarma"):
        self.base_dir = Path(base_dir)
        self.output_dir = self.base_dir / "mrtva_dolina"
        
        # Green screen settings
        self.green_color = np.array([0, 255, 0])
        self.tolerance = 30
        self.min_object_size = 400  # minimum 20x20 pixels
        
        # Source directories (prioritize green screen sources)
        self.sources_greenscreen = [
            self.base_dir / "Slike_za_Tiled",  # Has green screen
        ]
        
        self.sources_transparent = [
            self.base_dir / "assets" / "narezano_loceno",  # Already transparent
        ]
        
        # Category mapping (keywords → category)
        self.categories = {
            "characters/players": ["kai_character", "ana_character", "player"],
            "characters/npcs": ["npc", "ivan", "marija", "trader", "blacksmith", "merchant", "baker", "farmer", "doctor", "romance"],
            "characters/zombies": ["zombie"],
            "creatures/animals": ["animal", "cow", "chicken", "dog", "livestock", "donkey", "horse"],
            "creatures/monsters": ["troll", "creature", "anomalous", "cryptid", "slime"],
            "creatures/bosses": ["grok", "werewolf", "boss"],
            "buildings/houses": ["house", "farmhouse", "barn"],
            "buildings/ruins": ["ruin", "buildings_ruins"],
            "buildings/structures": ["building", "town_buildings", "structure", "blacksmith_workshop", "bridge", "wall", "city", "portal"],
            "terrain/ground": ["grass", "soil", "tileset", "ground", "terrain"],
            "terrain/water": ["water", "ocean", "pond", "river"],
            "terrain/special": ["mine_tileset", "mine_interior", "mine_entrance"],
            "vegetation/trees": ["tree", "sequoia", "birch", "vegetation"],
            "vegetation/plants": ["flower", "plant", "leaf"],
            "vegetation/crops": ["crop", "seed", "wheat", "vegetable", "fruit"],
            "objects/farming": ["sprinkler", "tool", "farming"],
            "objects/items": ["items_pack", "weapon", "bow", "arrow", "backpack", "chest", "grave"],
            "objects/furniture": ["furniture", "bathroom", "bedroom", "kitchen", "living_room"],
        }
        
        # Statistics
        self.stats = {
            "processed": 0,
            "separated": 0,
            "skipped": 0,
            "errors": 0
        }
        
        # Metadata collection
        self.metadata = []
    
    def create_output_structure(self):
        """Create all category directories"""
        print("🗂️  Creating directory structure...")
        for category in self.categories.keys():
            category_path = self.output_dir / category
            category_path.mkdir(parents=True, exist_ok=True)
            print(f"  ✓ {category}")
        
        # Also create misc category
        (self.output_dir / "misc").mkdir(parents=True, exist_ok=True)
        print(f"  ✓ misc")
    
    def is_green_screen(self, img_array):
        """Check if image has green screen background"""
        if img_array.shape[2] < 3:
            return False
        
        # Check for green pixels
        green_mask = np.all(np.abs(img_array[:, :, :3] - self.green_color) < self.tolerance, axis=2)
        green_percentage = np.sum(green_mask) / (img_array.shape[0] * img_array.shape[1])
        
        return green_percentage > 0.1  # At least 10% green
    
    def remove_green_screen(self, img):
        """Remove green screen and make transparent"""
        img_array = np.array(img.convert("RGBA"))
        
        # Create mask for green pixels
        green_mask = np.all(np.abs(img_array[:, :, :3] - self.green_color) < self.tolerance, axis=2)
        
        # Set alpha to 0 for green pixels
        img_array[green_mask, 3] = 0
        
        return Image.fromarray(img_array)
    
    def separate_objects(self, img):
        """Detect and separate individual objects"""
        img_array = np.array(img)
        
        # Get alpha channel (transparency)
        if img_array.shape[2] == 4:
            alpha = img_array[:, :, 3]
        else:
            # No alpha channel, use luminance threshold
            gray = np.mean(img_array[:, :, :3], axis=2)
            alpha = (gray > 10).astype(np.uint8) * 255
        
        # Label connected components
        labels, num_objects = ndimage.label(alpha > 128)
        
        if num_objects <= 1:
            return [img]  # Single object or empty
        
        # Extract each object
        objects = []
        for i in range(1, num_objects + 1):
            mask = (labels == i)
            
            # Check object size
            if np.sum(mask) < self.min_object_size:
                continue
            
            # Find bounding box
            rows = np.any(mask, axis=1)
            cols = np.any(mask, axis=0)
            
            if not np.any(rows) or not np.any(cols):
                continue
            
            y_min, y_max = np.where(rows)[0][[0, -1]]
            x_min, x_max = np.where(cols)[0][[0, -1]]
            
            # Crop object
            cropped = img_array[y_min:y_max+1, x_min:x_max+1]
            objects.append(Image.fromarray(cropped))
        
        return objects if objects else [img]
    
    def categorize_filename(self, filename):
        """Determine category from filename"""
        filename_lower = filename.lower()
        
        for category, keywords in self.categories.items():
            if any(keyword in filename_lower for keyword in keywords):
                return category
        
        return "misc"
    
    def process_image(self, image_path, test_mode=False, has_greenscreen=True):
        """Process a single image"""
        try:
            # Load image
            img = Image.open(image_path)
            
            if has_greenscreen:
                # Check and remove green screen
                img_array = np.array(img.convert("RGBA"))
                if not self.is_green_screen(img_array):
                    if test_mode:
                        print(f"  ⊘ {image_path.name} - No green screen, skipping")
                    self.stats["skipped"] += 1
                    return
                
                # Remove green screen
                img_transparent = self.remove_green_screen(img)
            else:
                # Already transparent, just convert to RGBA
                img_transparent = img.convert("RGBA")
            
            # Separate objects
            objects = self.separate_objects(img_transparent)
            
            # Categorize
            category = self.categorize_filename(image_path.stem)
            
            # Save objects
            if len(objects) == 1:
                # Single object
                output_path = self.output_dir / category / f"{image_path.stem}.png"
                objects[0].save(output_path, "PNG")
                
                if test_mode:
                    print(f"  ✓ {image_path.name} → {category}/{output_path.name}")
                
                self.metadata.append({
                    "filename": output_path.name,
                    "category": category,
                    "source": str(image_path),
                    "objects": 1
                })
                
            else:
                # Multiple objects
                for idx, obj in enumerate(objects, 1):
                    output_path = self.output_dir / category / f"{image_path.stem}_obj{idx:02d}.png"
                    obj.save(output_path, "PNG")
                
                if test_mode:
                    print(f"  ✓ {image_path.name} → {category}/ ({len(objects)} objects)")
                
                self.metadata.append({
                    "filename": f"{image_path.stem}_obj*.png",
                    "category": category,
                    "source": str(image_path),
                    "objects": len(objects)
                })
                
                self.stats["separated"] += len(objects)
            
            self.stats["processed"] += 1
            
        except Exception as e:
            print(f"  ✗ ERROR: {image_path.name} - {str(e)}")
            self.stats["errors"] += 1
    
    def test_run(self, num_samples=10):
        """Test processing on a few images"""
        print(f"\n🧪 TEST RUN - Processing {num_samples} sample images\n")
        
        # Collect sample images from GREENSCREEN sources
        samples = []
        for source in self.sources_greenscreen:
            if source.exists():
                png_files = list(source.glob("**/*.png"))[:num_samples]
                samples.extend(png_files)
        
        samples = samples[:num_samples]
        
        print(f"Found {len(samples)} test images from green screen sources\n")
        
        # Process samples
        for img_path in samples:
            self.process_image(img_path, test_mode=True, has_greenscreen=True)
        
        # Show stats
        print(f"\n📊 Test Results:")
        print(f"  Processed: {self.stats['processed']}")
        print(f"  Separated: {self.stats['separated']}")
        print(f"  Skipped: {self.stats['skipped']}")
        print(f"  Errors: {self.stats['errors']}")
        
        return self.stats["errors"] == 0
    
    def process_all(self):
        """Process all images"""
        print("\n🚀 PROCESSING ALL IMAGES\n")
        
        # Reset stats
        self.stats = {"processed": 0, "separated": 0, "skipped": 0, "errors": 0}
        
        # Process green screen images
        print("📗 Processing GREEN SCREEN images (Slike_za_Tiled)...")
        greenscreen_images = []
        for source in self.sources_greenscreen:
            if source.exists():
                png_files = list(source.glob("**/*.png"))
                greenscreen_images.extend(png_files)
                print(f"  Found {len(png_files)} images in {source.name}")
        
        for idx, img_path in enumerate(greenscreen_images, 1):
            if idx % 100 == 0:
                print(f"  Progress: {idx}/{len(greenscreen_images)} ({idx/len(greenscreen_images)*100:.1f}%)")
            self.process_image(img_path, has_greenscreen=True)
        
        print(f"  ✓ Processed {len(greenscreen_images)} green screen images\n")
        
        # Process transparent images (copy & categorize)
        print("📘 Processing TRANSPARENT images (narezano_loceno)...")
        transparent_images = []
        for source in self.sources_transparent:
            if source.exists():
                png_files = list(source.glob("**/*.png"))
                transparent_images.extend(png_files)
                print(f"  Found {len(png_files)} images in {source.name}")
        
        for idx, img_path in enumerate(transparent_images, 1):
            if idx % 100 == 0:
                print(f"  Progress: {idx}/{len(transparent_images)} ({idx/len(transparent_images)*100:.1f}%)")
            self.process_image(img_path, has_greenscreen=False)
        
        print(f"  ✓ Processed {len(transparent_images)} transparent images\n")
        
        total = len(greenscreen_images) + len(transparent_images)
        print(f"\n✅ DONE! Processed {total} images")
        print(f"\n📊 Final Statistics:")
        print(f"  Processed: {self.stats['processed']}")
        print(f"  Separated: {self.stats['separated']}")
        print(f"  Skipped: {self.stats['skipped']}")
        print(f"  Errors: {self.stats['errors']}")
        
        # Save metadata
        metadata_path = self.output_dir / "metadata.json"
        with open(metadata_path, 'w') as f:
            json.dump(self.metadata, f, indent=2)
        
        print(f"\n💾 Metadata saved to: {metadata_path}")

if __name__ == "__main__":
    print("=" * 60)
    print("  MASTER ASSET PROCESSOR - MRTVA DOLINA")
    print("=" * 60)
    
    processor = MrtvaDolina_AssetProcessor()
    
    # Create directory structure
    processor.create_output_structure()
    
    # Test run first
    print("\n" + "=" * 60)
    success = processor.test_run(num_samples=10)
    
    if success:
        print("\n✅ Test run successful!")
        print("\n" + "=" * 60)
        print("STARTING FULL PROCESSING...")
        print("=" * 60)
        
        # Run full processing
        processor.process_all()
        
        print("\n" + "=" * 60)
        print("🎉 ALL DONE!")
        print("=" * 60)
    else:
        print("\n❌ Test run had errors. Please fix before continuing.")