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

def process_image(image, target_aspect="9:16", fit_mode="crop"):
    """
    Frames an image to a target aspect ratio.
    fit_mode: 'crop' (center crop) or 'pad' (add borders)
    """
    if image is None:
        return None

    # Convert to PIL if needed (Gradio passes numpy arrays sometimes)
    if isinstance(image, np.ndarray):
        img = Image.fromarray(image)
    else:
        img = image

    w, h = img.size
    current_ratio = w / h

    # Parse target aspect
    if target_aspect == "9:16":
        target_ratio = 9 / 16
    elif target_aspect == "16:9":
        target_ratio = 16 / 9
    elif target_aspect == "1:1":
        target_ratio = 1.0
    elif target_aspect == "4:5":
        target_ratio = 4 / 5
    else:
        target_ratio = 9 / 16 # Default

    if fit_mode == "pad":
        # Calculate new size to fit inside target ratio
        if current_ratio > target_ratio:
            # Image is wider than target -> fit width, pad height
            new_w = w
            new_h = int(w / target_ratio)
        else:
            # Image is taller than target -> fit height, pad width
            new_h = h
            new_w = int(h * target_ratio)

        # Create background (black or blurred version could be added, using black for now)
        new_img = Image.new('RGB', (new_w, new_h), (0, 0, 0))

        # Paste original in center
        x_offset = (new_w - w) // 2
        y_offset = (new_h - h) // 2
        new_img.paste(img, (x_offset, y_offset))
        return new_img

    else: # Crop mode (Center Crop)
        # Calculate crop box
        if current_ratio > target_ratio:
            # Image is wider -> crop width
            new_w = int(h * target_ratio)
            new_h = h
            x_offset = (w - new_w) // 2
            y_offset = 0
        else:
            # Image is taller -> crop height
            new_w = w
            new_h = int(w / target_ratio)
            x_offset = 0
            y_offset = (h - new_h) // 2

        left = x_offset
        top = y_offset
        right = left + new_w
        bottom = top + new_h

        return img.crop((left, top, right, bottom))

with gr.Blocks(title="MMahBot Framing Studio") as demo:
    gr.Markdown("## 🖼️ AI Image Framer & Cropper")
    gr.Markdown("Prepare images for Watermarking and Video Generation by enforcing aspect ratios.")

    with gr.Row():
        with gr.Column():
            input_img = gr.Image(label="Upload Image", type="pil")
            aspect_radio = gr.Radio(
                ["9:16", "16:9", "1:1", "4:5"],
                value="9:16",
                label="Target Aspect Ratio"
            )
            mode_radio = gr.Radio(
                ["crop", "pad"],
                value="crop",
                label="Fitting Mode (Crop removes edges, Pad adds borders)"
            )
            btn = gr.Button("Process Frame", variant="primary")

        with gr.Column():
            output_img = gr.Image(label="Framed Result", type="pil")

    btn.click(fn=process_image, inputs=[input_img, aspect_radio, mode_radio], outputs=output_img)

if __name__ == "__main__":
    demo.launch()