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
        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."""
    padding_color = (255, 255, 255) if random.choice([True, False]) else (0, 0, 0)
    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:
        new_width = target_size[0]
        new_height = int(new_width / aspect_ratio)
    else:
        new_height = target_size[1]
        new_width = int(new_height * aspect_ratio)

    img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)

    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 = []

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

        resized_images = [resize_and_pad(img, (512, 512)) for img in group]
        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))

        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)

    # 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)

    zip_path = "combined_images.zip"
    shutil.make_archive("combined_images", 'zip', output_dir)
    return combined_images, zip_path

# Gradio UI
def main_interface():
    with gr.Blocks(css=".gradio-container { height: 100vh; }") as interface:
        with gr.Row():
            gr.Markdown("### Image Resizer and Combiner")

        with gr.Row():
            with gr.Column(scale=1, elem_id="left-panel"):
                uploaded_files = gr.Files(file_types=["image"], label="Upload Images", interactive=True)
                combine_button = gr.Button("Combine and Process", elem_id="combine-button")

            with gr.Column(scale=2, elem_id="right-panel"):
                preview_gallery = gr.Gallery(label="Preview Combined Images", columns=3, elem_id="preview-gallery")
                download_zip_button = gr.Button("Download ZIP", elem_id="download-zip")

        # Button actions
        combine_button.click(
            process_images,
            inputs=[uploaded_files],
            outputs=[preview_gallery, download_zip_button],
        )

        download_zip_button.click(
            lambda x: x,
            inputs=[],
            outputs=[download_zip_button],
        )

    return interface

main_interface().launch()