#!/usr/bin/env python3
"""
🎨 Background Removal Script - Batch Transparency
Removes white (#FFFFFF) and black (#000000) backgrounds from images
Creates clean PNG-32 with alpha channel

Special handling for Kai & Gronk sprites - sharp edges, no white haze!
"""

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

# Configuration
WHITE_THRESHOLD = 250  # Pixels with RGB all above this = white
BLACK_THRESHOLD = 5    # Pixels with RGB all below this = black
EDGE_TOLERANCE = 30    # For edge detection near colored areas

def remove_background(image_path, output_path=None):
    """
    Remove white and black backgrounds from an image.
    Creates transparent PNG-32.
    """
    try:
        img = Image.open(image_path)
        
        # Convert to RGBA if needed
        if img.mode != 'RGBA':
            img = img.convert('RGBA')
        
        # Get image data as numpy array
        data = np.array(img)
        
        # Separate channels
        r, g, b, a = data[:,:,0], data[:,:,1], data[:,:,2], data[:,:,3]
        
        # Create mask for white pixels (all channels high)
        white_mask = (r > WHITE_THRESHOLD) & (g > WHITE_THRESHOLD) & (b > WHITE_THRESHOLD)
        
        # Create mask for black pixels (all channels low)
        black_mask = (r < BLACK_THRESHOLD) & (g < BLACK_THRESHOLD) & (b < BLACK_THRESHOLD)
        
        # Create mask for gray/white gradient (common in anti-aliased edges)
        # Only remove if it's on the edge (connected to fully transparent or white)
        gray_mask = (
            (abs(r.astype(int) - g.astype(int)) < 10) & 
            (abs(g.astype(int) - b.astype(int)) < 10) & 
            (abs(r.astype(int) - b.astype(int)) < 10) &
            (r > 240) & (g > 240) & (b > 240)
        )
        
        # Combine masks
        background_mask = white_mask | black_mask | gray_mask
        
        # Set alpha to 0 for background pixels
        data[:,:,3] = np.where(background_mask, 0, a)
        
        # Create new image
        result = Image.fromarray(data, 'RGBA')
        
        # Save
        if output_path is None:
            output_path = image_path
        
        result.save(output_path, 'PNG', optimize=True)
        return True
        
    except Exception as e:
        print(f"  ❌ Error processing {image_path}: {e}")
        return False

def remove_background_advanced(image_path, output_path=None):
    """
    Advanced background removal with edge-aware processing.
    Preserves sharp edges on colored areas (like Kai's dreads).
    """
    try:
        img = Image.open(image_path)
        
        if img.mode != 'RGBA':
            img = img.convert('RGBA')
        
        data = np.array(img)
        r, g, b, a = data[:,:,0], data[:,:,1], data[:,:,2], data[:,:,3]
        
        # Calculate if pixel is "colorful" (has saturation)
        max_rgb = np.maximum(np.maximum(r, g), b)
        min_rgb = np.minimum(np.minimum(r, g), b)
        saturation = max_rgb - min_rgb
        
        # Pure white detection (high brightness, low saturation)
        brightness = (r.astype(int) + g.astype(int) + b.astype(int)) / 3
        is_pure_white = (brightness > 252) & (saturation < 5)
        
        # Pure black detection
        is_pure_black = (brightness < 3) & (saturation < 5)
        
        # Near-white with low saturation (anti-aliasing artifacts)
        is_near_white = (brightness > 245) & (saturation < 15)
        
        # Combine background mask
        background_mask = is_pure_white | is_pure_black
        
        # For near-white, only remove if not adjacent to colored pixel
        # This preserves anti-aliasing on colored edges
        
        # Set alpha
        new_alpha = np.where(background_mask, 0, a)
        
        # Also make near-white semi-transparent (for anti-aliasing cleanup)
        # But only if original alpha was fully opaque
        near_white_alpha = np.where(
            is_near_white & (a > 250),
            np.maximum(0, 255 - brightness.astype(int)),
            new_alpha
        )
        
        data[:,:,3] = near_white_alpha.astype(np.uint8)
        
        result = Image.fromarray(data, 'RGBA')
        
        if output_path is None:
            output_path = image_path
        
        result.save(output_path, 'PNG', optimize=True)
        return True
        
    except Exception as e:
        print(f"  ❌ Error processing {image_path}: {e}")
        return False

def process_directory(directory, use_advanced=True):
    """
    Process all PNG images in a directory recursively.
    """
    dir_path = Path(directory)
    if not dir_path.exists():
        print(f"❌ Directory not found: {directory}")
        return 0, 0
    
    extensions = ['*.png', '*.PNG']
    all_images = []
    
    for ext in extensions:
        all_images.extend(dir_path.rglob(ext))
    
    all_images = list(set(all_images))
    total = len(all_images)
    success = 0
    
    print(f"\n📁 Processing {directory}")
    print(f"   Found {total} PNG images")
    
    for i, img_path in enumerate(all_images):
        img_str = str(img_path)
        name = img_path.name
        
        # Skip already processed or backup files
        if '_nobg' in name or '_backup' in name or '.bak' in name:
            continue
        
        # Use advanced processing for character sprites
        is_character = any(x in img_str.lower() for x in ['kai', 'gronk', 'ana', 'susi', 'character'])
        
        if use_advanced or is_character:
            result = remove_background_advanced(img_str)
        else:
            result = remove_background(img_str)
        
        if result:
            success += 1
            print(f"   ✅ [{i+1}/{total}] {name}")
        else:
            print(f"   ❌ [{i+1}/{total}] {name} - FAILED")
    
    return success, total

def main():
    print("=" * 60)
    print("🎨 BATCH BACKGROUND REMOVAL - NovaFarma Assets")
    print("=" * 60)
    print("\nRemoving white (#FFFFFF) and black (#000000) backgrounds")
    print("Creating transparent PNG-32 with alpha channel")
    print("Special handling for Kai & Gronk sprites!\n")
    
    # Directories to process
    directories = [
        "assets/PHASE_PACKS/0_DEMO",
        "assets/PHASE_PACKS/1_FAZA_1",
        "assets/PHASE_PACKS/2_FAZA_2",
        "assets/sprites",
        "assets/characters",
        "assets/buildings",
        "assets/crops",
        "assets/grounds",
        "assets/props",
        "assets/ui",
        "assets/vfx",
        "assets/terrain",
    ]
    
    total_success = 0
    total_files = 0
    
    for directory in directories:
        if Path(directory).exists():
            success, total = process_directory(directory)
            total_success += success
            total_files += total
    
    print("\n" + "=" * 60)
    print(f"✅ COMPLETED!")
    print(f"   Processed: {total_success}/{total_files} images")
    print(f"   Failed: {total_files - total_success}")
    print("=" * 60)

if __name__ == '__main__':
    main()