#!/usr/bin/env python3
"""
Advanced background removal for game assets:
- Detects background color (white, black, or other)
- Preserves bold cartoon outlines
- Creates clean transparent PNGs
- Works with dual-style system
"""

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

def detect_background_color(img):
    """
    Detect dominant background color by sampling corners
    Returns: 'white', 'black', 'color', or None
    """
    width, height = img.size
    
    # Sample corner pixels (10x10 from each corner)
    corners = [
        img.crop((0, 0, 10, 10)),                      # Top-left
        img.crop((width-10, 0, width, 10)),            # Top-right
        img.crop((0, height-10, 10, height)),          # Bottom-left
        img.crop((width-10, height-10, width, height)) # Bottom-right
    ]
    
    # Get average color from corners
    corner_colors = []
    for corner in corners:
        corner_array = np.array(corner)
        avg_color = corner_array.mean(axis=(0, 1))
        corner_colors.append(avg_color)
    
    avg_bg = np.mean(corner_colors, axis=0)
    
    # Determine background type
    if avg_bg[:3].mean() > 240:  # Very light
        return 'white'
    elif avg_bg[:3].mean() < 15:  # Very dark
        return 'black'
    else:
        return 'color'

def remove_white_background(img, tolerance=30):
    """
    Remove white/light background while preserving cartoon outlines
    """
    # Convert to RGBA
    img = img.convert('RGBA')
    data = np.array(img)
    
    r, g, b, a = data[:,:,0], data[:,:,1], data[:,:,2], data[:,:,3]
    
    # Create mask: white areas (all channels high)
    white_mask = (r > 255 - tolerance) & (g > 255 - tolerance) & (b > 255 - tolerance)
    
    # Create mask: light gray areas
    light_gray_mask = (
        (r > 220) & (g > 220) & (b > 220) &
        (np.abs(r.astype(int) - g.astype(int)) < tolerance) &
        (np.abs(r.astype(int) - b.astype(int)) < tolerance) &
        (np.abs(g.astype(int) - b.astype(int)) < tolerance)
    )
    
    # Combine masks
    bg_mask = white_mask | light_gray_mask
    
    # Set alpha to 0 where background detected
    data[bg_mask, 3] = 0
    
    return Image.fromarray(data, mode='RGBA')

def remove_black_background(img, tolerance=60):
    """
    Remove black/dark background while preserving black outlines
    More aggressive version to catch dark gray artifacts
    """
    # Convert to RGBA
    img = img.convert('RGBA')
    data = np.array(img)
    
    r, g, b, a = data[:,:,0], data[:,:,1], data[:,:,2], data[:,:,3]
    
    # Detect edges (where color changes rapidly = outlines!)
    from scipy import ndimage
    
    # Calculate gradients (more sensitive)
    grad_r = ndimage.sobel(r)
    grad_g = ndimage.sobel(g)
    grad_b = ndimage.sobel(b)
    
    # Magnitude of gradient
    magnitude = np.sqrt(grad_r**2 + grad_g**2 + grad_b**2)
    
    # Areas with high gradient = edges (keep these!)
    # Lowered threshold to preserve more edges
    is_edge = magnitude > 30
    
    # Create mask: dark areas (all channels low) AND not an edge
    # More aggressive: also catch dark gray (60 instead of 30)
    black_mask = (
        (r < tolerance) & 
        (g < tolerance) & 
        (b < tolerance) & 
        ~is_edge  # NOT an edge!
    )
    
    # Set alpha to 0 where background detected
    data[black_mask, 3] = 0
    
    return Image.fromarray(data, mode='RGBA')

def remove_color_background(img, sample_size=10, tolerance=30):
    """
    Remove colored background by sampling corners
    """
    # Convert to RGBA
    img = img.convert('RGBA')
    data = np.array(img)
    width, height = img.size
    
    # Sample corner to get BG color
    corner = img.crop((0, 0, sample_size, sample_size))
    bg_color = np.array(corner).mean(axis=(0, 1))[:3]
    
    r, g, b, a = data[:,:,0], data[:,:,1], data[:,:,2], data[:,:,3]
    
    # Create mask: colors similar to background
    color_diff = np.sqrt(
        (r - bg_color[0])**2 + 
        (g - bg_color[1])**2 + 
        (b - bg_color[2])**2
    )
    
    bg_mask = color_diff < tolerance
    
    # Set alpha to 0 where background detected
    data[bg_mask, 3] = 0
    
    return Image.fromarray(data, mode='RGBA')

def cleanup_edges(img):
    """
    Clean up semi-transparent edge pixels (anti-aliasing artifacts)
    + Add subtle feathering for smoother edges
    """
    from scipy import ndimage
    
    data = np.array(img)
    alpha = data[:,:,3].astype(float)
    
    # Slightly blur the alpha channel for smooth edges
    alpha_blurred = ndimage.gaussian_filter(alpha, sigma=0.5)
    
    # Remove very transparent pixels (< 10% opacity)
    alpha_blurred[alpha_blurred < 25] = 0
    
    # Make mostly opaque pixels fully opaque (> 90%)
    alpha_blurred[alpha_blurred > 230] = 255
    
    data[:,:,3] = alpha_blurred.astype(np.uint8)
    
    return Image.fromarray(data, mode='RGBA')

def remove_background_auto(img_path: Path, output_path: Path = None, dry_run: bool = False):
    """
    Automatically detect and remove background
    """
    try:
        # Load image
        img = Image.open(img_path)
        
        # Skip if already has transparency
        if img.mode == 'RGBA':
            data = np.array(img)
            alpha = data[:,:,3]
            # If already has transparent pixels, skip
            if (alpha < 255).sum() > (alpha.size * 0.1):  # >10% transparent
                print(f"  ⏭️  Already has transparency: {img_path.name}")
                return False
        
        # Detect background type
        bg_type = detect_background_color(img)
        
        if bg_type is None:
            print(f"  ⚠️  Could not detect background: {img_path.name}")
            return False
        
        print(f"  🎨 Detected {bg_type} background: {img_path.name}")
        
        # Remove background based on type
        if bg_type == 'white':
            result = remove_white_background(img, tolerance=30)
        elif bg_type == 'black':
            result = remove_black_background(img, tolerance=60)  # More aggressive!
        else:
            result = remove_color_background(img, tolerance=30)
        
        # Cleanup edges
        result = cleanup_edges(result)
        
        # Save
        if output_path is None:
            output_path = img_path
        
        if not dry_run:
            result.save(output_path, 'PNG', optimize=True)
            print(f"  ✅ Removed {bg_type} background: {img_path.name}")
        
        return True
        
    except Exception as e:
        print(f"  ❌ Error processing {img_path.name}: {e}")
        return False

def process_directory(directory: Path, dry_run: bool = False, recursive: bool = True):
    """
    Process all PNG files in directory
    """
    print(f"\n{'='*70}")
    print(f"📂 Processing: {directory}")
    print(f"{'='*70}")
    
    if recursive:
        png_files = list(directory.rglob('*.png'))
    else:
        png_files = list(directory.glob('*.png'))
    
    if not png_files:
        print("  No PNG files found")
        return 0
    
    count = 0
    for png_file in png_files:
        if remove_background_auto(png_file, dry_run=dry_run):
            count += 1
    
    return count

def main():
    import argparse
    
    parser = argparse.ArgumentParser(description='Advanced background removal for game assets')
    parser.add_argument('path', type=str, help='File or directory to process')
    parser.add_argument('--dry-run', action='store_true', help='Show what would be done')
    parser.add_argument('--no-recursive', action='store_true', help='Do not process subdirectories')
    args = parser.parse_args()
    
    print("=" * 70)
    print("🧹 ADVANCED BACKGROUND REMOVAL")
    print("=" * 70)
    
    if args.dry_run:
        print("\n⚠️  DRY RUN MODE - No changes will be made\n")
    
    print("\nFeatures:")
    print("  • Auto-detects white, black, or colored backgrounds")
    print("  • Preserves cartoon outlines and edges")
    print("  • Creates clean transparent PNGs")
    print("=" * 70)
    
    path = Path(args.path)
    
    if path.is_file():
        remove_background_auto(path, dry_run=args.dry_run)
    elif path.is_dir():
        total = process_directory(path, dry_run=args.dry_run, recursive=not args.no_recursive)
        print(f"\n{'='*70}")
        if not args.dry_run:
            print(f"✅ COMPLETE! Processed {total} images")
        else:
            print(f"✅ DRY RUN COMPLETE! Would process {total} images")
        print("=" * 70)
    else:
        print(f"❌ Error: {path} is not a valid file or directory")

if __name__ == "__main__":
    main()