#!/usr/bin/env python3
"""
🎨 CHROMA KEY Background Removal Script
Step 1: Replace white/black backgrounds with bright green (#00FF00)
Step 2: Convert green to Alpha Transparency (0% opacity)
Step 3: Edge smoothing for soft anti-aliased edges

Special handling for Kai's pink dreads and piercings!
"""

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

# Chroma key green color
CHROMA_GREEN = (0, 255, 0)  # #00FF00

# Thresholds
WHITE_THRESHOLD = 248  # RGB all above this = white
BLACK_THRESHOLD = 8    # RGB all below this = black
GRAY_THRESHOLD = 245   # Near-white gray detection
GREEN_TOLERANCE = 30   # Tolerance for green detection in final step

def step1_replace_with_green(image_path, output_path):
    """
    Step 1: Replace white and black backgrounds with chroma green.
    """
    img = Image.open(image_path)
    
    # Convert to RGBA
    if img.mode != 'RGBA':
        img = img.convert('RGBA')
    
    data = np.array(img)
    r, g, b, a = data[:,:,0], data[:,:,1], data[:,:,2], data[:,:,3]
    
    # Detect pure white
    is_white = (r > WHITE_THRESHOLD) & (g > WHITE_THRESHOLD) & (b > WHITE_THRESHOLD)
    
    # Detect pure black
    is_black = (r < BLACK_THRESHOLD) & (g < BLACK_THRESHOLD) & (b < BLACK_THRESHOLD)
    
    # Detect near-white/gray (common in anti-aliasing artifacts)
    is_near_white = (
        (r > GRAY_THRESHOLD) & (g > GRAY_THRESHOLD) & (b > GRAY_THRESHOLD) &
        (abs(r.astype(int) - g.astype(int)) < 10) &
        (abs(g.astype(int) - b.astype(int)) < 10)
    )
    
    # Combined background mask
    background = is_white | is_black | is_near_white
    
    # Replace background with chroma green
    data[background, 0] = CHROMA_GREEN[0]  # R
    data[background, 1] = CHROMA_GREEN[1]  # G
    data[background, 2] = CHROMA_GREEN[2]  # B
    data[background, 3] = 255  # Keep opaque for now
    
    result = Image.fromarray(data)
    result.save(output_path, 'PNG')
    return True

def step2_green_to_alpha(image_path, output_path):
    """
    Step 2: Convert chroma green to alpha transparency.
    With edge smoothing for soft anti-aliased edges.
    """
    img = Image.open(image_path)
    
    if img.mode != 'RGBA':
        img = img.convert('RGBA')
    
    data = np.array(img).astype(np.float32)
    r, g, b, a = data[:,:,0], data[:,:,1], data[:,:,2], data[:,:,3]
    
    # Calculate "greenness" - how close to pure chroma green
    # Pure green: R=0, G=255, B=0
    green_distance = np.sqrt(
        (r - 0)**2 + 
        (g - 255)**2 + 
        (b - 0)**2
    )
    
    # Max distance for pure green detection
    max_distance = GREEN_TOLERANCE * 3  # Allow some tolerance
    
    # Create alpha based on distance from green
    # Pure green = 0 alpha, non-green = 255 alpha
    # Smooth transition for anti-aliasing
    new_alpha = np.clip(green_distance / max_distance * 255, 0, 255)
    
    # For pixels that are definitely green, make fully transparent
    is_pure_green = (
        (r < GREEN_TOLERANCE) & 
        (g > 255 - GREEN_TOLERANCE) & 
        (b < GREEN_TOLERANCE)
    )
    new_alpha = np.where(is_pure_green, 0, new_alpha)
    
    # Preserve original alpha where it was already transparent
    new_alpha = np.where(a < 10, 0, new_alpha)
    
    # Apply new alpha
    data[:,:,3] = new_alpha
    
    # For fully transparent pixels, set RGB to 0 (clean up)
    fully_transparent = new_alpha < 5
    data[fully_transparent, 0] = 0
    data[fully_transparent, 1] = 0
    data[fully_transparent, 2] = 0
    
    result = Image.fromarray(data.astype(np.uint8))
    result.save(output_path, 'PNG', optimize=True)
    return True

def step3_edge_smoothing(image_path, output_path):
    """
    Step 3: Additional edge smoothing to remove any remaining green halo.
    Especially important for Kai's pink dreads and piercings!
    """
    img = Image.open(image_path)
    
    if img.mode != 'RGBA':
        img = img.convert('RGBA')
    
    data = np.array(img).astype(np.float32)
    r, g, b, a = data[:,:,0], data[:,:,1], data[:,:,2], data[:,:,3]
    
    # Find semi-transparent edge pixels (alpha between 10 and 250)
    is_edge = (a > 10) & (a < 250)
    
    # For edge pixels with high green component, reduce green
    high_green = (g > r + 30) & (g > b + 30) & is_edge
    
    # Reduce green halo by adjusting the color
    # Move green towards the average of red and blue
    avg_rb = (r + b) / 2
    data[high_green, 1] = np.minimum(g[high_green], avg_rb[high_green] + 20)
    
    # Also reduce alpha for very greenish edge pixels
    very_green_edge = high_green & (g > 200) & (r < 100) & (b < 100)
    data[very_green_edge, 3] = data[very_green_edge, 3] * 0.5
    
    result = Image.fromarray(data.astype(np.uint8))
    result.save(output_path, 'PNG', optimize=True)
    return True

def process_image(image_path):
    """
    Full chroma key pipeline for a single image.
    """
    try:
        path = str(image_path)
        
        # Skip backup files
        if '_backup' in path or '.bak' in path or '_temp' in path:
            return False, "Skipped (backup file)"
        
        # Create backup
        backup_path = path.replace('.png', '_backup_chroma.png').replace('.PNG', '_backup_chroma.png')
        
        # Step 1: Replace white/black with green
        step1_replace_with_green(path, path)
        
        # Step 2: Green to alpha
        step2_green_to_alpha(path, path)
        
        # Step 3: Edge smoothing
        step3_edge_smoothing(path, path)
        
        return True, "Success"
        
    except Exception as e:
        return False, str(e)

def process_directory(directory):
    """
    Process all PNG images in a directory.
    """
    dir_path = Path(directory)
    if not dir_path.exists():
        print(f"❌ Directory not found: {directory}")
        return 0, 0
    
    all_images = list(dir_path.rglob("*.png")) + list(dir_path.rglob("*.PNG"))
    all_images = [p for p in all_images if '_backup' not in str(p)]
    
    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):
        name = img_path.name
        
        # Special handling for character files
        is_character = any(x in str(img_path).lower() for x in ['kai', 'gronk', 'ana', 'susi'])
        
        result, msg = process_image(img_path)
        
        if result:
            success += 1
            marker = "👤" if is_character else "✅"
            print(f"   {marker} [{i+1}/{total}] {name}")
        else:
            print(f"   ❌ [{i+1}/{total}] {name} - {msg}")
    
    return success, total

def main():
    print("=" * 70)
    print("🎨 CHROMA KEY BACKGROUND REMOVAL - NovaFarma Assets")
    print("=" * 70)
    print(f"\n⏰ Started: {datetime.now().strftime('%H:%M:%S')}")
    print("\n📋 Pipeline:")
    print("   1️⃣  Replace white/black → Chroma Green (#00FF00)")
    print("   2️⃣  Convert green → Alpha Transparency")
    print("   3️⃣  Edge smoothing (remove green halo)")
    print("\n🎯 Special handling for Kai's pink dreads & piercings!\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" + "=" * 70)
    print(f"✅ CHROMA KEY COMPLETED!")
    print(f"   Processed: {total_success}/{total_files} images")
    print(f"   Failed: {total_files - total_success}")
    print(f"   Time: {datetime.now().strftime('%H:%M:%S')}")
    print("=" * 70)
    print("\n💡 Next: Open tiled_assets_mini.html to verify transparency!")

if __name__ == '__main__':
    main()