#!/usr/bin/env python3
"""
AUTOMATIC BACKGROUND REMOVAL SYSTEM
Removes green backgrounds from all reference images and makes them transparent.
Uses color-based removal (green #00FF00 background).
"""

import os
from PIL import Image
import numpy as np

def remove_green_background(input_path, output_path):
    """
    Remove green (#00FF00) background and make transparent.
    """
    print(f"Processing: {input_path}")
    
    # Open image
    img = Image.open(input_path).convert("RGBA")
    data = np.array(img)
    
    # Get RGB channels
    r, g, b, a = data.T
    
    # Define green background color range (accounting for compression artifacts)
    # Pure green is #00FF00 (0, 255, 0)
    green_threshold = 200  # Flexibility for compression
    
    # Find green pixels
    green_areas = (g > green_threshold) & (r < 100) & (b < 100)
    
    # Make green pixels transparent
    data[..., 3][green_areas.T] = 0
    
    # Save with transparency
    result = Image.fromarray(data)
    result.save(output_path, 'PNG')
    print(f"✅ Saved: {output_path}")

def process_references_folder(references_path):
    """
    Process all PNG/JPG images in references folder recursively.
    """
    processed_count = 0
    
    for root, dirs, files in os.walk(references_path):
        for file in files:
            # Only process image files
            if file.lower().endswith(('.png', '.jpg', '.jpeg')):
                input_path = os.path.join(root, file)
                
                # Skip if already processed (has '_nobg' in name)
                if '_nobg' in file:
                    continue
                
                # Create output path (same location, add _nobg suffix)
                base_name = os.path.splitext(file)[0]
                output_path = os.path.join(root, f"{base_name}_nobg.png")
                
                try:
                    remove_green_background(input_path, output_path)
                    processed_count += 1
                except Exception as e:
                    print(f"❌ Error processing {input_path}: {e}")
    
    return processed_count

if __name__ == "__main__":
    references_path = os.path.join(os.path.dirname(__file__), '..', 'references')
    
    print("=" * 60)
    print("🎨 AUTOMATIC BACKGROUND REMOVAL")
    print("=" * 60)
    print(f"Processing folder: {references_path}")
    print()
    
    count = process_references_folder(references_path)
    
    print()
    print("=" * 60)
    print(f"✅ COMPLETE! Processed {count} images")
    print("=" * 60)