#!/usr/bin/env python3
"""
🟢 BATCH GREEN SCREEN GENERATOR
Converts ALL images to have #00FF00 chroma key background
Uses rembg AI for clean subject detection

Usage: python3 scripts/batch_green_screen_all.py
"""

import os
import sys
from pathlib import Path
from rembg import remove
from PIL import Image
from datetime import datetime
import shutil

# 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/crops",
]

CHROMA_GREEN = (0, 255, 0, 255)  # #00FF00
BACKUP_DIR = "assets/_backup_before_greenscreen"

def process_image(input_path, create_backup=True):
    """
    Process single image:
    1. Create backup of original
    2. Use AI to remove background
    3. Add green #00FF00 background
    4. Overwrite original with green version
    """
    try:
        # Open image
        img = Image.open(input_path)
        
        # Create backup if requested
        if create_backup:
            backup_path = input_path.replace("assets/", f"{BACKUP_DIR}/")
            os.makedirs(os.path.dirname(backup_path), exist_ok=True)
            shutil.copy(input_path, backup_path)
        
        # Step 1: AI remove background
        img_no_bg = remove(img)
        
        # Step 2: Create green background
        green_bg = Image.new('RGBA', img_no_bg.size, CHROMA_GREEN)
        
        # Step 3: Composite subject onto green
        green_bg.paste(img_no_bg, (0, 0), img_no_bg)
        
        # Step 4: Convert to RGB and save (overwrite original)
        green_rgb = green_bg.convert('RGB')
        green_rgb.save(input_path, 'PNG')
        
        return True, "OK"
        
    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
    
    # Find all PNG files
    all_images = list(dir_path.rglob("*.png")) + list(dir_path.rglob("*.PNG"))
    
    # Filter out backups and already processed
    all_images = [p for p in all_images if '_backup' not in str(p) and '_GREEN' not in str(p)]
    
    total = len(all_images)
    success = 0
    
    print(f"\n📁 {directory}")
    print(f"   Found {total} PNG images")
    
    for i, img_path in enumerate(all_images):
        name = img_path.name
        result, msg = process_image(str(img_path))
        
        if result:
            success += 1
            print(f"   🟢 [{i+1}/{total}] {name}")
        else:
            print(f"   ❌ [{i+1}/{total}] {name} - {msg}")
    
    return success, total

def main():
    print("=" * 60)
    print("🟢 BATCH GREEN SCREEN GENERATOR")
    print("=" * 60)
    print(f"\n⏰ Started: {datetime.now().strftime('%H:%M:%S')}")
    print(f"🎯 Target: Convert all images to #00FF00 green background")
    print(f"💾 Backups: Saved to {BACKUP_DIR}/")
    print("\n⚠️  This will OVERWRITE original files!")
    print("    Backups are created automatically.\n")
    
    # Create backup directory
    os.makedirs(BACKUP_DIR, exist_ok=True)
    
    total_success = 0
    total_files = 0
    
    for directory in DIRECTORIES:
        success, total = process_directory(directory)
        total_success += success
        total_files += total
    
    print("\n" + "=" * 60)
    print(f"✅ COMPLETED!")
    print(f"   Processed: {total_success}/{total_files} images")
    print(f"   Failed: {total_files - total_success}")
    print(f"   Backups: {BACKUP_DIR}/")
    print(f"   Time: {datetime.now().strftime('%H:%M:%S')}")
    print("=" * 60)
    
    # Instructions for reverting
    print("\n📋 To REVERT to originals:")
    print(f"   cp -r {BACKUP_DIR}/* assets/")

if __name__ == '__main__':
    main()