#!/usr/bin/env python3
"""
Background Removal for DolinaSmrti Assets
Uses rembg to remove white backgrounds and create transparent PNGs
"""

import os
import sys
from pathlib import Path
from PIL import Image

try:
    from rembg import remove
except ImportError:
    print("❌ rembg not installed. Installing now...")
    os.system("pip3 install rembg")
    from rembg import remove


def remove_background(input_path: Path, output_path: Path = None):
    """Remove background from image and save as transparent PNG"""
    
    if output_path is None:
        output_path = input_path
    
    print(f"  🎨 Processing: {input_path.name}")
    
    # Load image
    with open(input_path, 'rb') as input_file:
        input_data = input_file.read()
    
    # Remove background
    output_data = remove(input_data)
    
    # Save as PNG
    with open(output_path, 'wb') as output_file:
        output_file.write(output_data)
    
    print(f"  ✅ Saved: {output_path.name}")


def process_directory(directory: Path, recursive: bool = False):
    """Process all PNG files in directory"""
    
    print(f"\n{'='*60}")
    print(f"📁 Processing directory: {directory}")
    print(f"{'='*60}")
    
    pattern = "**/*.png" if recursive else "*.png"
    png_files = list(directory.glob(pattern))
    
    total = len(png_files)
    print(f"\nFound {total} PNG files")
    
    for i, png_file in enumerate(png_files, 1):
        print(f"\n[{i}/{total}]")
        remove_background(png_file)
    
    print(f"\n{'='*60}")
    print(f"✅ Completed! Processed {total} images")
    print(f"{'='*60}")


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python3 remove_background.py <image_or_directory>")
        print("\nExamples:")
        print("  python3 remove_background.py image.png")
        print("  python3 remove_background.py assets/slike/zgradbe/")
        print("  python3 remove_background.py assets/images/ --recursive")
        sys.exit(1)
    
    path = Path(sys.argv[1])
    recursive = "--recursive" in sys.argv or "-r" in sys.argv
    
    if not path.exists():
        print(f"❌ Path not found: {path}")
        sys.exit(1)
    
    if path.is_file():
        remove_background(path)
    elif path.is_dir():
        process_directory(path, recursive=recursive)
    else:
        print(f"❌ Invalid path: {path}")
        sys.exit(1)