import gradio as gr
from PIL import Image
import os

# Function to process and display the images in grid format
def create_collage(uploaded_files, grid_layout):
    # Open images
    images = [Image.open(file.name) for file in uploaded_files]

    # Determine grid dimensions based on layout
    if grid_layout == "1x2":
        rows, cols = 1, 2
    elif grid_layout == "1x3":
        rows, cols = 1, 3
    elif grid_layout == "2x2":
        rows, cols = 2, 2
    elif grid_layout == "2x3":
        rows, cols = 2, 3

    # Create a blank canvas for the collage without spacing
    collage_width = max(img.width for img in images) * cols
    collage_height = max(img.height for img in images) * rows
    collage = Image.new("RGB", (collage_width, collage_height))

    # Resize images to fit the grid
    resized_images = []
    for img in images:
        img_width = collage_width // cols
        img_height = collage_height // rows
        img_resized = img.resize((img_width, img_height), Image.Resampling.LANCZOS)
        resized_images.append(img_resized)

    # Place images on the canvas in grid format without spacing
    for i, img in enumerate(resized_images):
        row = i // cols
        col = i % cols
        collage.paste(img, (col * img.width, row * img.height))

    # Save the collage as a PNG file in a temporary directory
    output_path = "temp_collage.png"
    collage.save(output_path, format="PNG")

    # Return the file path for download
    return output_path

# Gradio Interface
with gr.Blocks() as app:
    with gr.Column():
        # Step 1: Title and description
        gr.Markdown("# **Image Collage Maker**")
        gr.Markdown("Create custom image collages with different grid layouts like 1x2, 1x3, 2x2, and 2x3. Upload your images and select a layout!")

        # Step 2: File upload for multiple images
        uploaded_files = gr.Files(label="Upload Images", file_count="multiple", file_types=["image"])

        # Step 3: Grid selection
        grid_layout = gr.Radio(
            choices=["1x2", "1x3", "2x2", "2x3"],
            label="Select Grid Layout",
            value="1x2"  # Default value
        )

        # Step 4: Display preview of the grid
        output = gr.Image(label="Collage Preview")

        # Step 5: Process the images when files are uploaded or grid layout is selected
        uploaded_files.change(create_collage, inputs=[uploaded_files, grid_layout], outputs=output)
        grid_layout.change(create_collage, inputs=[uploaded_files, grid_layout], outputs=output)

        # Step 6: Download button for the final collage as PNG
        download_button = gr.Button("Download Collage (PNG)")
        download_button.click(create_collage, inputs=[uploaded_files, grid_layout], outputs=gr.File(label="Download Collage"))

app.launch(share=True)  # Public link to share the app