#!/usr/bin/env python3
"""
Smart Background Removal Script
- Preserves character details (faces, hands)
- Only removes solid backgrounds (white, green, checkerboard)
- Uses tolerance range to avoid removing character colors
"""

from PIL import Image
import os
import sys

def smart_remove_background(image_path, output_path=None):
    """
    Odstrani samo čista ozadja, ohrani detajle like
    """
    if output_path is None:
        output_path = image_path
    
    img = Image.open(image_path).convert("RGBA")
    width, height = img.size
    pixels = img.load()
    
    # Definiramo barve ozadij za odstranitev
    backgrounds_to_remove = [
        # Pure White
        ((245, 245, 245), (255, 255, 255)),
        # Pure Green (Chroma Key)
        ((0, 245, 0), (10, 255, 10)),
        # Light Gray (Checkerboard)
        ((200, 200, 200), (220, 220, 220)),
    ]
    
    # Varnostna območja - NE BRIŠI teh barv (kože toni)
    skin_tones = [
        # Light skin
        ((240, 210, 180), (255, 230, 210)),
        # Medium skin  
        ((210, 180, 140), (240, 210, 180)),
        # Dark skin
        ((140, 100, 70), (210, 180, 140)),
    ]
    
    removed_count = 0
    
    for y in range(height):
        for x in range(width):
            r, g, b, a = pixels[x, y]
            
            # Preveri, če je to kožni ton - OHRANI
            is_skin = False
            for (min_r, min_g, min_b), (max_r, max_g, max_b) in skin_tones:
                if min_r <= r <= max_r and min_g <= g <= max_g and min_b <= b <= max_b:
                    is_skin = True
                    break
            
            if is_skin:
                continue  # Ohrani piksel
            
            # Preveri, če je to ozadje za odstranitev
            for (min_r, min_g, min_b), (max_r, max_g, max_b) in backgrounds_to_remove:
                if min_r <= r <= max_r and min_g <= g <= max_g and min_b <= b <= max_b:
                    pixels[x, y] = (r, g, b, 0)  # Naredi transparentno
                    removed_count += 1
                    break
    
    img.save(output_path, "PNG")
    print(f"✅ {os.path.basename(output_path)}: Removed {removed_count} background pixels")
    return removed_count > 0


def process_directory(directory):
    """Process all PNG images in directory"""
    processed = 0
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.lower().endswith('.png'):
                file_path = os.path.join(root, file)
                try:
                    if smart_remove_background(file_path):
                        processed += 1
                except Exception as e:
                    print(f"❌ Error processing {file}: {e}")
    
    print(f"\n✅ Total processed: {processed} images")


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python3 smart_background_removal.py <directory>")
        sys.exit(1)
    
    target_dir = sys.argv[1]
    if not os.path.exists(target_dir):
        print(f"❌ Directory not found: {target_dir}")
        sys.exit(1)
    
    print(f"🔍 Processing images in: {target_dir}")
    process_directory(target_dir)
    print("✨ Done!")