import os
import sys
import subprocess

# Define the source directory on the Desktop
SOURCE_DIR = "/Users/davidkotnik/Desktop/referencne slike 2"

def install_dependencies():
    """Attempt to install required packages if they are missing."""
    packages = ["rembg", "Pillow", "onnxruntime"]
    for package in packages:
        try:
            __import__(package)
        except ImportError:
            if package == "Pillow":
                 # Pillow is imported as PIL
                try:
                    __import__("PIL")
                    continue
                except ImportError:
                    pass
            
            print(f"Installing {package}...")
            try:
                subprocess.check_call([sys.executable, "-m", "pip", "install", package])
            except subprocess.CalledProcessError as e:
                print(f"Failed to install {package}: {e}")
                print("Please try installing it manually: pip install rembg Pillow onnxruntime")
                sys.exit(1)

def process_images():
    # Verify directory
    if not os.path.exists(SOURCE_DIR):
        print(f"Error: Directory {SOURCE_DIR} not found.")
        return

    # Import libraries AFTER installation
    try:
        from rembg import remove
        from PIL import Image
    except ImportError:
        print("Error: Library import failed even after installation attempt.")
        return

    # Get PNG files
    files = [f for f in os.listdir(SOURCE_DIR) if f.lower().endswith('.png')]
    total = len(files)
    
    if total == 0:
        print("No PNG files found to process.")
        return

    print(f"Found {total} PNG images. Starting AI background removal...")
    print("WARNING: This makes permanent changes to the PNG files (JPG backups are safe).")
    print("This might take a while depending on your computer speed.")

    success = 0
    errors = 0

    for i, filename in enumerate(files):
        file_path = os.path.join(SOURCE_DIR, filename)
        
        try:
            print(f"[{i+1}/{total}] Processing: {filename}")
            
            # Open the image
            with open(file_path, "rb") as input_file:
                input_data = input_file.read()
            
            # Remove background
            output_data = remove(input_data)
            
            # Save the result back to the same file (overwrite)
            with open(file_path, "wb") as output_file:
                output_file.write(output_data)
                
            success += 1
            
        except Exception as e:
            print(f"❌ Error processing {filename}: {e}")
            errors += 1

    print(f"\n--- DONE ---")
    print(f"Successful: {success}")
    print(f"Errors: {errors}")

if __name__ == "__main__":
    install_dependencies()
    process_images()