#!/usr/bin/env python3
"""
Sprite Cleanup & Optimization Tool for Novafarma
Fixes transparency issues and optimizes file sizes

Features:
- Removes checkerboard backgrounds
- Converts to proper transparent PNG
- Optimizes file size (500KB+ → 50-100KB)
- Batch processing all sprites

Usage:
    python sprite_cleanup.py
"""

from PIL import Image
import os
import glob

def remove_checkerboard_background(image):
    """
    Remove checkerboard pattern (common in AI-generated images)
    Target: Gray/white checkerboard (R,G,B values between 150-255 and similar)
    """
    img = image.convert('RGBA')
    data = img.getdata()
    
    new_data = []
    for item in data:
        r, g, b, a = item
        
        # Detect checkerboard gray/white
        # Checkerboard is usually perfect gray (R≈G≈B) with high values
        is_gray = abs(r - g) < 30 and abs(g - b) <30
        is_light = r > 150 and g > 150 and b > 150
        
        if is_gray and is_light:
            # Make transparent
            new_data.append((r, g, b, 0))
        else:
            # Keep original
            new_data.append(item)
    
    img.putdata(new_data)
    return img

def remove_white_background(image):
    """Remove pure white or near-white backgrounds"""
    img = image.convert('RGBA')
    data = img.getdata()
    
    new_data = []
    for item in data:
        r, g, b, a = item
        
        # Near-white threshold
        if r > 240 and g > 240 and b > 240:
            new_data.append((r, g, b, 0))
        else:
            new_data.append(item)
    
    img.putdata(new_data)
    return img

def remove_solid_color_background(image, target_color=(200, 200, 200), tolerance=30):
    """Remove specific solid color background with tolerance"""
    img = image.convert('RGBA')
    data = img.getdata()
    
    target_r, target_g, target_b = target_color
    new_data = []
    
    for item in data:
        r, g, b, a = item
        
        # Check if color is within tolerance of target
        if (abs(r - target_r) < tolerance and 
            abs(g - target_g) < tolerance and 
            abs(b - target_b) < tolerance):
            new_data.append((r, g, b, 0))
        else:
            new_data.append(item)
    
    img.putdata(new_data)
    return img

def optimize_sprite(input_file, output_file=None, aggressive=True):
    """
    Optimize sprite:
    1. Remove backgrounds
    2. Reduce file size
    3. Save as optimized PNG
    """
    
    print(f"📦 Processing: {input_file}")
    
    if not os.path.exists(input_file):
        print(f"❌ File not found: {input_file}")
        return False
    
    # Original file size
    original_size = os.path.getsize(input_file) / 1024  # KB
    
    try:
        # Load image
        img = Image.open(input_file)
        
        # Get original dimensions
        width, height = img.size
        
        # Remove backgrounds (multiple passes for best results)
        img = remove_checkerboard_background(img)
        img = remove_white_background(img)
        
        # Optional: Resize if too large (AI images are often huge)
        if aggressive and (width > 512 or height > 512):
            # Calculate new size (max 512px on longest side)
            ratio = min(512 / width, 512 / height)
            new_width = int(width * ratio)
            new_height = int(height * ratio)
            
            print(f"   ⚙️ Resizing: {width}x{height} → {new_width}x{new_height}")
            img = img.resize((new_width, new_height), Image.LANCZOS)
        
        # Output filename
        if output_file is None:
            # Create backup
            backup_file = input_file.replace('.png', '_backup.png')
            if not os.path.exists(backup_file):
                os.rename(input_file, backup_file)
                print(f"   💾 Backup saved: {backup_file}")
            output_file = input_file
        
        # Save optimized
        img.save(output_file, 'PNG', optimize=True, compress_level=9)
        
        # New file size
        new_size = os.path.getsize(output_file) / 1024  # KB
        reduction = ((original_size - new_size) / original_size) * 100
        
        print(f"   ✅ Saved: {output_file}")
        print(f"   📊 Size: {original_size:.1f}KB → {new_size:.1f}KB ({reduction:.1f}% reduction)")
        
        return True
        
    except Exception as e:
        print(f"   ❌ Error: {e}")
        return False

def batch_optimize_directory(directory, pattern='*.png', aggressive=True):
    """Optimize all PNG files in a directory"""
    
    files = glob.glob(os.path.join(directory, pattern))
    
    if not files:
        print(f"No files found matching {pattern} in {directory}")
        return
    
    print(f"\n🎨 Optimizing {len(files)} files in {directory}")
    print("=" * 60)
    
    success_count = 0
    total_original_size = 0
    total_new_size = 0
    
    for file in files:
        # Skip backups
        if '_backup' in file or '_strip' in file:
            print(f"⏭️ Skipping: {os.path.basename(file)}")
            continue
        
        original_size = os.path.getsize(file) / 1024
        
        if optimize_sprite(file, aggressive=aggressive):
            success_count += 1
            new_size = os.path.getsize(file) / 1024
            total_original_size += original_size
            total_new_size += new_size
        
        print()  # Blank line between files
    
    # Summary
    print("=" * 60)
    print(f"✅ Successfully processed: {success_count}/{len(files)} files")
    if total_original_size > 0:
        total_reduction = ((total_original_size - total_new_size) / total_original_size) * 100
        print(f"📊 Total size: {total_original_size:.1f}KB → {total_new_size:.1f}KB")
        print(f"💾 Space saved: {total_original_size - total_new_size:.1f}KB ({total_reduction:.1f}%)")

def create_optimized_player_sprite():
    """Create a simple, clean player sprite (placeholder)"""
    
    print("\n🎨 Creating optimized player sprite...")
    
    # 64x64 pixel art player
    size = 64
    img = Image.new('RGBA', (size, size), (0, 0, 0, 0))
    pixels = img.load()
    
    # Simple character (head + body + legs)
    # Head (circle)
    for y in range(20, 35):
        for x in range(25, 40):
            dx = x - 32
            dy = y - 27
            if dx*dx + dy*dy < 64:  # Circle radius
                pixels[x, y] = (255, 200, 150, 255)  # Skin tone
    
    # Body (rectangle)
    for y in range(35, 50):
        for x in range(27, 37):
            pixels[x, y] = (100, 150, 200, 255)  # Blue shirt
    
    # Legs (two rectangles)
    for y in range(50, 64):
        for x in range(27, 32):
            pixels[x, y] = (50, 50, 100, 255)  # Dark pants
        for x in range(32, 37):
            pixels[x, y] = (50, 50, 100, 255)
    
    output_file = 'player_sprite_clean.png'
    img.save(output_file, 'PNG', optimize=True)
    print(f"✅ Created: {output_file}")
    
    return output_file

if __name__ == '__main__':
    print("🎮 Novafarma Sprite Cleanup Tool\n")
    
    # Option 1: Clean all sprites in assets folder
    assets_dir = 'c:/novafarma/assets'
    
    if os.path.exists(assets_dir):
        print("🔧 Mode: Aggressive optimization (resize + cleanup)")
        print("⚠️  Backups will be created automatically\n")
        
        # Process main assets folder
        batch_optimize_directory(assets_dir, '*.png', aggressive=True)
        
        # Process sprites subfolder
        sprites_dir = os.path.join(assets_dir, 'sprites')
        if os.path.exists(sprites_dir):
            batch_optimize_directory(sprites_dir, '*.png', aggressive=False)  # Don't resize sprite sheets
    
    else:
        print(f"❌ Assets directory not found: {assets_dir}")
        print("\n💡 Usage:")
        print("   1. Update assets_dir path in script")
        print("   2. Run: python sprite_cleanup.py")
        print("\n   Or use individual sprite:")
        print("   optimize_sprite('path/to/sprite.png')")
    
    print("\n✨ Done! Check your assets folder for optimized sprites.")
    print("📁 Backups are saved as *_backup.png")