import gradio as gr
from PIL import Image, ImageOps
import numpy as np
import cv2


def get_new_size_and_padding(old_w, old_h, target_ratio):
    old_ratio = old_w / old_h
    if old_ratio > target_ratio:
        # Image is wider than target: pad height
        new_w = old_w
        new_h = int(old_w / target_ratio)
        pad_top = (new_h - old_h) // 2
        pad_bottom = new_h - old_h - pad_top
        pad_left = pad_right = 0
    else:
        # Image is taller than target: pad width
        new_h = old_h
        new_w = int(old_h * target_ratio)
        pad_left = (new_w - old_w) // 2
        pad_right = new_w - old_w - pad_left
        pad_top = pad_bottom = 0
    return (new_w, new_h), (pad_left, pad_top, pad_right, pad_bottom)


def image_filler(img, padding):
    # Use OpenCV to blur the border region for a simple "image filler"
    np_img = np.array(img)
    h, w = np_img.shape[:2]
    pad_left, pad_top, pad_right, pad_bottom = padding

    # Create a blurred version of the image
    blurred = cv2.GaussianBlur(np_img, (51, 51), 0)

    # Create a new image with the blurred background
    new_h = h + pad_top + pad_bottom
    new_w = w + pad_left + pad_right
    result = np.zeros((new_h, new_w, 3), dtype=np.uint8)
    result[:, :] = blurred[0, 0]  # Fill with a color from the image

    # Place the blurred image in the background
    result[:, :] = blurred[0, 0]
    result[pad_top : pad_top + h, pad_left : pad_left + w] = np_img

    # Blend the edges for a smooth transition
    # (For simplicity, just use the blurred background)
    return Image.fromarray(result)


def resize_with_border(image, aspect_ratio, border_fill):
    # Parse aspect ratio
    w_ratio, h_ratio = map(int, aspect_ratio.split(":"))
    target_ratio = w_ratio / h_ratio

    # Ensure image is RGB
    image = image.convert("RGB")
    old_w, old_h = image.size

    (new_w, new_h), padding = get_new_size_and_padding(old_w, old_h, target_ratio)

    if border_fill == "White":
        # Use Pillow's expand with white fill
        result = ImageOps.expand(image, border=padding, fill=(255, 255, 255))
    else:
        # Use image filler (blurred border)
        result = image_filler(image, padding)

    # Resize to exact target size (optional, in case of rounding)
    result = result.resize((new_w, new_h), Image.LANCZOS)
    return result


aspect_ratios = ["1:1", "4:3", "16:9", "3:2", "21:9", "9:16"]

with gr.Blocks() as demo:
    gr.Markdown("# Image Resizer with Border Fill")
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Upload Image")
            aspect_input = gr.Dropdown(
                aspect_ratios, value="1:1", label="Target Aspect Ratio"
            )
            fill_input = gr.Radio(
                ["White", "Image Filler"], value="White", label="Border Fill"
            )
            submit_btn = gr.Button("Resize Image")
        with gr.Column():
            image_output = gr.Image(type="pil", label="Resized Image")
    submit_btn.click(
        resize_with_border,
        inputs=[image_input, aspect_input, fill_input],
        outputs=image_output,
    )

demo.launch()