from PIL import Image
import os

# Lista slik za procesiranje
images_to_process = [
    'assets/tree_purple.png',
    'assets/tree_apple.png',
    'assets/tree_pear.png',
    'assets/tree_cherry.png'
]

def remove_background(image_path):
    """Odstrani ozadje iz slike - naredi transparentno"""
    try:
        # Naloži sliko
        img = Image.open(image_path).convert('RGBA')
        data = img.getdata()
        
        new_data = []
        for item in data:
            r, g, b, a = item
            
            # Izračunaj brightness in grayscale
            brightness = (r + g + b) / 3
            is_grayscale = abs(r - g) < 20 and abs(g - b) < 20 and abs(r - b) < 20
            
            # Odstrani:
            # 1. Vse sive barve (šahovsko ozadje)
            # 2. Vse svetle barve (brightness > 150)
            # 3. Bele barve
            # 4. Pastelne barve (nizka saturacija)
            
            should_remove = False
            
            # Sive barve
            if is_grayscale and brightness > 80:
                should_remove = True
            
            # Svetle barve
            if brightness > 150:
                should_remove = True
            
            # Bele
            if r > 240 and g > 240 and b > 240:
                should_remove = True
            
            # Pastelne (nizka saturacija)
            max_rgb = max(r, g, b)
            min_rgb = min(r, g, b)
            saturation = 0 if max_rgb == 0 else (max_rgb - min_rgb) / max_rgb
            
            if saturation < 0.3 and brightness > 120:
                should_remove = True
            
            if should_remove:
                new_data.append((r, g, b, 0))  # Transparentno
            else:
                new_data.append(item)  # Ohrani original
        
        # Posodobi sliko
        img.putdata(new_data)
        
        # Shrani
        img.save(image_path)
        print(f'✅ Processed: {image_path}')
        
    except Exception as e:
        print(f'❌ Error processing {image_path}: {e}')

# Procesiraj vse slike
print('🎨 Removing backgrounds from tree sprites...')
for img_path in images_to_process:
    if os.path.exists(img_path):
        remove_background(img_path)
    else:
        print(f'⚠️ File not found: {img_path}')

print('✅ Done! All backgrounds removed.')