#!/usr/bin/env python3
"""
POPOLNO ČIŠČENJE VIZUALNIH VIROV
- Pretvori vse JPG → PNG z chroma key
- Odstrani zeleno ozadje
- Standardiziraj velikosti (trees 512px, camp/props 256px)
- Pobriši stare JPG datoteke
"""

import os
from PIL import Image
import numpy as np

# Direktoriji za čiščenje
TARGET_DIRS = [
    "assets/slike/nova mapa faza 0-1",
    "assets/slike/glavna_referenca"
]

# Chroma key threshold
GREEN_THRESHOLD = 100  # sensitivity for green detection

def remove_green_background(image):
    """Odstrani zeleno ozadje in naredi prosojno"""
    img = image.convert("RGBA")
    data = np.array(img)
    
    # RGB channels
    r, g, b, a = data.T
    
    # Zelena areas: G > R and G > B with some threshold
    green_areas = (g > r + GREEN_THRESHOLD) & (g > b + GREEN_THRESHOLD)
    
    # Also remove very light/white areas (common in greenscreen)
    white_areas = (r > 240) & (g > 240) & (b > 240)
    
    # Set alpha to 0 for green and white areas
    data[..., 3][green_areas.T] = 0
    data[..., 3][white_areas.T] = 0
    
    return Image.fromarray(data)

def is_tree_or_large_prop(filename):
    """Preveri če je datoteka drevo ali velik prop (potrebuje 512px)"""
    tree_keywords = ['drevo', 'tree', 'wood', 'log', 'stump']
    return any(keyword in filename.lower() for keyword in tree_keywords)

def process_jpg_to_png(jpg_path, output_dir):
    """Proces JPG → PNG z chroma key"""
    try:
        # Naloži sliko
        img = Image.open(jpg_path)
        
        # Odstrani zeleno ozadje
        img_transparent = remove_green_background(img)
        
        # Določi target velikost
        filename = os.path.basename(jpg_path)
        target_size = 512 if is_tree_or_large_prop(filename) else 256
        
        # Resize ohranjujoč aspect ratio
        img_transparent.thumbnail((target_size, target_size), Image.Resampling.LANCZOS)
        
        # Ustvari PNG ime
        png_filename = os.path.splitext(filename)[0] + ".png"
        png_path = os.path.join(output_dir, png_filename)
        
        # Shrani PNG
        img_transparent.save(png_path, "PNG")
        print(f"✅ {filename} → {png_filename} ({target_size}px)")
        
        return png_path
        
    except Exception as e:
        print(f"❌ Error processing {jpg_path}: {e}")
        return None

def clean_existing_pngs(png_dir):
    """Očisti obstoječe PNG datoteke ki majo še zeleno ozadje"""
    png_files = []
    for root, dirs, files in os.walk(png_dir):
        for file in files:
            if file.lower().endswith('.png') and not file.startswith('.'):
                png_files.append(os.path.join(root, file))
    
    cleaned_count = 0
    for png_path in png_files:
        try:
            img = Image.open(png_path)
            
            # Preveri če ima že alpha channel
            if img.mode != 'RGBA':
                img = img.convert('RGBA')
            
            # Odstrani zeleno
            img_cleaned = remove_green_background(img)
            
            # Shrani nazaj
            img_cleaned.save(png_path, "PNG")
            cleaned_count += 1
            print(f"🧹 Cleaned: {os.path.basename(png_path)}")
            
        except Exception as e:
            print(f"⚠️  Could not clean {png_path}: {e}")
    
    return cleaned_count

def main():
    print("=" * 60)
    print("🧹 POPOLNO ČIŠČENJE VIZUALNIH VIROV")
    print("=" * 60)
    
    total_converted = 0
    total_deleted = 0
    total_cleaned = 0
    
    for target_dir in TARGET_DIRS:
        if not os.path.exists(target_dir):
            print(f"⚠️  Directory not found: {target_dir}")
            continue
        
        print(f"\n📁 Processing: {target_dir}")
        print("-" * 60)
        
        # Find all JPG files
        jpg_files = []
        for root, dirs, files in os.walk(target_dir):
            for file in files:
                if file.lower().endswith(('.jpg', '.jpeg')) and not file.startswith('.'):
                    jpg_files.append(os.path.join(root, file))
        
        print(f"Found {len(jpg_files)} JPG files to convert")
        
        # Convert JPG → PNG
        for jpg_path in jpg_files:
            output_dir = os.path.dirname(jpg_path)
            png_path = process_jpg_to_png(jpg_path, output_dir)
            
            if png_path:
                total_converted += 1
                # Delete original JPG
                try:
                    os.remove(jpg_path)
                    total_deleted += 1
                    print(f"🗑️  Deleted: {os.path.basename(jpg_path)}")
                except Exception as e:
                    print(f"⚠️  Could not delete {jpg_path}: {e}")
        
        # Clean existing PNGs
        print(f"\n🧹 Cleaning existing PNG files in {target_dir}...")
        cleaned = clean_existing_pngs(target_dir)
        total_cleaned += cleaned
    
    # Summary
    print("\n" + "=" * 60)
    print("📊 SUMMARY")
    print("=" * 60)
    print(f"✅ JPG → PNG converted: {total_converted}")
    print(f"🗑️  JPG files deleted: {total_deleted}")
    print(f"🧹 PNG files cleaned: {total_cleaned}")
    print("=" * 60)
    print("✨ ČIŠČENJE KONČANO!")
    print("=" * 60)

if __name__ == "__main__":
    main()