import gradio as gr
from PIL import Image, ImageOps
import numpy as np
import os
import uuid

# Ensure there's a directory for outputs
os.makedirs("outputs", exist_ok=True)

def make_square(img, size=3000, fill_color=(0, 0, 0)):
    x, y = img.size
    scale = size / max(x, y)
    new_size = (int(x * scale), int(y * scale))
    # Replace deprecated ANTIALIAS with modern equivalent
    img = img.resize(new_size, Image.Resampling.LANCZOS)
    new_img = Image.new("RGB", (size, size), fill_color)
    new_img.paste(img, ((size - new_size[0]) // 2, (size - new_size[1]) // 2))
    return new_img

def blend_images(images):
    if len(images) < 2:
        return "Upload at least two images.", None
    
    try:
        # Add error handling for image processing
        processed = []
        for img in images:
            try:
                processed.append(make_square(Image.open(img)))
            except Exception as e:
                return f"Error processing image: {str(e)}", None
                
        base = np.array(processed[0]).astype(np.float32)

        for img in processed[1:]:
            base = (base + np.array(img).astype(np.float32)) / 2

        final = Image.fromarray(np.uint8(base))

        # Save to file
        output_path = f"outputs/amalgam_{uuid.uuid4().hex[:8]}.png"
        final.save(output_path)

        return final, output_path
    except Exception as e:
        return f"Error during blending: {str(e)}", None

demo = gr.Interface(
    fn=blend_images,
    inputs=gr.File(file_types=["image"], file_count="multiple", label="Upload 2–5 stills"),
    outputs=[
        gr.Image(label="Blended Image"),
        gr.File(label="Download Image")
    ],
    title="Amalgamator",
    description="Upload up to 5 stills. Outputs a 3000x3000 blended image preserving the aesthetic. Save it as PNG below."
)

demo.launch()