#!/usr/bin/env python3
"""
BATCH VISUAL ASSET CLEANUP
Removes backgrounds from ALL image assets in the project.
Processes: /references/, /assets/, /style_test_samples/

Features:
- Multi-folder scanning
- Chroma keying (green #00FF00 removal)
- Smart detection of already transparent images
- Progress tracking
- Automatic _nobg.png generation
"""

import os
from PIL import Image
import numpy as np
from pathlib import Path

# Configuration
SEARCH_FOLDERS = [
    'references',
    'assets/images',
    'assets/bugs',
    'assets/tools',
    'assets/crops',
    'assets/buildings',
    'assets/characters',
    'assets/enemies',
    'assets/npcs',
    'style_test_samples'
]

GREEN_KEY = (0, 255, 0)  # #00FF00
TOLERANCE = 30  # Color variance tolerance

def has_transparency(image):
    """Check if image already has transparent pixels"""
    if image.mode != 'RGBA':
        return False
    
    alpha = np.array(image)[:, :, 3]
    return np.any(alpha < 255)

def remove_green_background(image_path, output_path):
    """Remove green chroma key background and save as transparent PNG"""
    try:
        img = Image.open(image_path).convert('RGBA')
        data = np.array(img)
        
        # Check if already transparent
        if has_transparency(img):
            print(f"  ⏭️  Already transparent: {os.path.basename(image_path)}")
            return False
        
        # Extract RGB channels
        red, green, blue, alpha = data[:, :, 0], data[:, :, 1], data[:, :, 2], data[:, :, 3]
        
        # Create mask for green pixels (with tolerance)
        mask = (
            (np.abs(red.astype(int) - GREEN_KEY[0]) <= TOLERANCE) &
            (np.abs(green.astype(int) - GREEN_KEY[1]) <= TOLERANCE) &
            (np.abs(blue.astype(int) - GREEN_KEY[2]) <= TOLERANCE)
        )
        
        # Check if green background exists
        green_pixel_count = np.sum(mask)
        if green_pixel_count == 0:
            print(f"  ℹ️  No green background: {os.path.basename(image_path)}")
            return False
        
        # Set alpha to 0 for green pixels
        data[:, :, 3] = np.where(mask, 0, 255)
        
        # Save as transparent PNG
        result = Image.fromarray(data, 'RGBA')
        result.save(output_path, 'PNG')
        
        percentage = (green_pixel_count / (data.shape[0] * data.shape[1])) * 100
        print(f"  ✅ Processed: {os.path.basename(image_path)} ({percentage:.1f}% background removed)")
        return True
        
    except Exception as e:
        print(f"  ❌ ERROR: {os.path.basename(image_path)} - {e}")
        return False

def scan_and_process(base_dir):
    """Scan all configured folders and process images"""
    stats = {
        'total_found': 0,
        'processed': 0,
        'skipped_transparent': 0,
        'skipped_no_green': 0,
        'errors': 0
    }
    
    print("🔍 SCANNING FOR IMAGES...\n")
    
    for folder in SEARCH_FOLDERS:
        folder_path = os.path.join(base_dir, folder)
        
        if not os.path.exists(folder_path):
            print(f"⏭️  Skipping (not found): {folder}")
            continue
        
        print(f"\n📁 Processing: {folder}/")
        print("=" * 60)
        
        # Find all image files
        image_files = []
        for ext in ['*.png', '*.jpg', '*.jpeg']:
            image_files.extend(Path(folder_path).rglob(ext))
        
        folder_stats = {
            'found': len(image_files),
            'processed': 0
        }
        
        for img_path in sorted(image_files):
            # Skip already processed files
            if '_nobg' in img_path.stem:
                continue
            
            stats['total_found'] += 1
            
            # Generate output filename
            output_path = img_path.parent / f"{img_path.stem}_nobg.png"
            
            # Skip if already exists
            if output_path.exists():
                print(f"  ⏭️  Already exists: {output_path.name}")
                stats['skipped_transparent'] += 1
                continue
            
            # Process the image
            result = remove_green_background(str(img_path), str(output_path))
            
            if result:
                stats['processed'] += 1
                folder_stats['processed'] += 1
            elif result is False:
                # Check reason from function output
                stats['skipped_no_green'] += 1
        
        print(f"\n📊 Folder Summary: {folder_stats['processed']}/{folder_stats['found']} processed")
    
    return stats

def main():
    """Main execution"""
    print("=" * 60)
    print("🎨 BATCH VISUAL ASSET CLEANUP")
    print("=" * 60)
    print(f"Target folders: {len(SEARCH_FOLDERS)}")
    print(f"Chroma key: RGB{GREEN_KEY} ±{TOLERANCE}")
    print("=" * 60)
    
    # Get project root
    script_dir = os.path.dirname(os.path.abspath(__file__))
    project_root = os.path.dirname(script_dir)
    
    # Process all folders
    stats = scan_and_process(project_root)
    
    # Final report
    print("\n" + "=" * 60)
    print("📊 FINAL REPORT")
    print("=" * 60)
    print(f"Total images found:        {stats['total_found']}")
    print(f"✅ Successfully processed:  {stats['processed']}")
    print(f"⏭️  Already transparent:     {stats['skipped_transparent']}")
    print(f"ℹ️  No green background:    {stats['skipped_no_green']}")
    print(f"❌ Errors:                  {stats['errors']}")
    print("=" * 60)
    
    if stats['processed'] > 0:
        print(f"\n✅ CLEANUP COMPLETE! {stats['processed']} backgrounds removed.")
    else:
        print("\n✅ All images already clean!")

if __name__ == "__main__":
    main()