import gradio as gr
from PIL import Image
import os
import shutil
import random

def process_transparent_image(image, padding_color):
    """Handle images with transparency by converting transparent areas to a solid color."""
    if image.mode in ("RGBA", "LA") or (image.mode == "P" and "transparency" in image.info):
        alpha = image.convert("RGBA").getchannel("A")  # Extract the alpha channel
        # Create a new image with the background color
        background = Image.new("RGBA", image.size, padding_color + (255,))
        background.paste(image, mask=alpha)  # Paste with transparency mask
        return background.convert("RGB")  # Convert back to RGB
    return image.convert("RGB")  # For non-transparent images

def resize_and_pad(image, target_size):
    """Resize the image and pad it to match target size."""
    # Randomly select black or white as the background color
    padding_color = (255, 255, 255) if random.choice([True, False]) else (0, 0, 0)

    # Handle transparent images
    img = process_transparent_image(image, padding_color)

    aspect_ratio = img.width / img.height
    target_aspect_ratio = target_size[0] / target_size[1]

    if aspect_ratio > target_aspect_ratio:
        # Wider than target aspect ratio
        new_width = target_size[0]
        new_height = int(new_width / aspect_ratio)
    else:
        # Taller than target aspect ratio
        new_height = target_size[1]
        new_width = int(new_height * aspect_ratio)

    # Use high-quality resampling for resizing
    img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)

    # Create a new image with the determined background color
    new_img = Image.new("RGB", target_size, padding_color)
    offset = ((target_size[0] - new_width) // 2, (target_size[1] - new_height) // 2)
    new_img.paste(img, offset)
    return new_img

def combine_images(images, target_size=(2048, 2048)):
    """Combine four images into a single one."""
    num_images = len(images)
    output_images = []

    # Group images into chunks of 4
    for i in range(0, num_images, 4):
        group = images[i:i+4]
        if len(group) < 4:
            continue

        # Resize and pad each image
        resized_images = [resize_and_pad(img, (512, 512)) for img in group]

        # Create a blank canvas
        combined_img = Image.new("RGB", (1024, 1024), (0, 0, 0))
        combined_img.paste(resized_images[0], (0, 0))
        combined_img.paste(resized_images[1], (512, 0))
        combined_img.paste(resized_images[2], (0, 512))
        combined_img.paste(resized_images[3], (512, 512))

        # Resize if the final size exceeds target
        if combined_img.width > target_size[0] or combined_img.height > target_size[1]:
            combined_img = combined_img.resize(target_size, Image.Resampling.LANCZOS)

        output_images.append(combined_img)

    return output_images

def process_images(uploaded_images):
    """Main processing function."""
    images = [Image.open(img) for img in uploaded_images]
    combined_images = combine_images(images)
    result_paths = []

    # Save combined images
    output_dir = "output_images"
    os.makedirs(output_dir, exist_ok=True)
    for idx, img in enumerate(combined_images):
        output_path = os.path.join(output_dir, f"combined_{idx + 1}.png")
        img.save(output_path)
        result_paths.append(output_path)

    # Create a ZIP file
    zip_path = "combined_images.zip"
    shutil.make_archive("combined_images", 'zip', output_dir)

    return zip_path

# Gradio UI
interface = gr.Interface(
    fn=process_images,
    inputs=gr.Files(file_types=["image"], label="Upload Images"),
    outputs=gr.File(label="Download ZIP"),
    title="Image Resizer and Combiner",
    description="Upload multiple images. This tool will resize and pad them to match a consistent size, then combine every four images into one. The output will be resized if it exceeds 2048x2048, and all results are available as a ZIP file."
)

interface.launch()