#!/usr/bin/env python3
"""
Standalone Background Removal Tool
Remove backgrounds from generated FM face images
"""

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

def remove_background_simple(image_path, output_path=None):
    """
    Simple background removal using PIL and basic image processing
    This is a fallback method when rembg is not available
    """
    try:
        with Image.open(image_path) as img:
            # Convert to RGBA if not already
            if img.mode != 'RGBA':
                img = img.convert('RGBA')

            # Get image data
            data = img.getdata()

            # Simple background removal based on color similarity
            # This is a basic implementation - results may vary
            new_data = []
            for item in data:
                # Consider pixels with high brightness/low saturation as background
                r, g, b, a = item

                # Calculate brightness and saturation
                brightness = (r + g + b) / 3
                saturation = max(r, g, b) - min(r, g, b)

                # If pixel is very bright and low saturation (likely background)
                if brightness > 240 and saturation < 30:
                    # Make transparent
                    new_data.append((r, g, b, 0))
                else:
                    new_data.append(item)

            # Create new image with transparent background
            new_img = Image.new('RGBA', img.size)
            new_img.putdata(new_data)

            # Save result
            output_path = output_path or image_path
            new_img.save(output_path, 'PNG')
            return True

    except Exception as e:
        print(f"Error processing {image_path}: {e}")
        return False

def find_face_images(directory):
    """Find all PNG images in directory that might be face images"""
    image_extensions = {'.png', '.jpg', '.jpeg'}
    face_images = []

    for file_path in Path(directory).rglob('*'):
        if file_path.suffix.lower() in image_extensions:
            # Skip if already processed (has _no_bg in name)
            if '_no_bg' not in file_path.stem:
                face_images.append(str(file_path))

    return sorted(face_images)

def main():
    parser = argparse.ArgumentParser(description="Remove backgrounds from FM face images")
    parser.add_argument("directory", help="Directory containing face images")
    parser.add_argument("--method", choices=['rembg', 'simple', 'auto'],
                       default='auto', help="Background removal method")
    parser.add_argument("--output-dir", help="Output directory for processed images")
    parser.add_argument("--suffix", default="_no_bg",
                       help="Suffix to add to processed images")

    args = parser.parse_args()

    if not os.path.exists(args.directory):
        print(f"Error: Directory {args.directory} does not exist")
        sys.exit(1)

    # Find images
    images = find_face_images(args.directory)
    if not images:
        print(f"No images found in {args.directory}")
        sys.exit(1)

    print(f"Found {len(images)} images to process")

    # Determine method
    method = args.method
    if method == 'auto':
        try:
            from rembg import remove
            method = 'rembg'
            print("Using rembg for background removal")
        except ImportError:
            method = 'simple'
            print("Using simple background removal (rembg not available)")

    # Process images
    success_count = 0
    output_dir = args.output_dir or args.directory

    for i, image_path in enumerate(images, 1):
        print(f"Processing {i}/{len(images)}: {os.path.basename(image_path)}")

        if method == 'rembg':
            try:
                with Image.open(image_path) as img:
                    output = remove(img)
                    output_path = os.path.join(output_dir,
                                             f"{Path(image_path).stem}{args.suffix}.png")
                    output.save(output_path, 'PNG')
                    success_count += 1
            except Exception as e:
                print(f"  Error: {e}")
        else:
            # Simple method
            output_path = os.path.join(output_dir,
                                     f"{Path(image_path).stem}{args.suffix}.png")
            if remove_background_simple(image_path, output_path):
                success_count += 1

    print(f"\nCompleted: {success_count}/{len(images)} images processed successfully")

    if success_count < len(images):
        print(f"\nNote: {len(images) - success_count} images failed to process")
        print("You may need to install rembg for better results:")
        print("pip install rembg[gpu]==2.0.59")

if __name__ == "__main__":
    main()