import gradio as gr
from gradio.components.image_editor import EditorValue
from gradio_imageslider import ImageSlider
from PIL import Image
from typing import cast
import numpy as np
from simple_lama_inpainting import SimpleLama


simple_lama = SimpleLama()

def HWC3(x):
    if x.ndim == 2:
        x = x[:, :, None]
    H, W, C = x.shape
    if C == 3:
        return x
    if C == 1:
        return np.concatenate([x, x, x], axis=2)
    if C == 4:
        color = x[:, :, 0:3].astype(np.float32)
        alpha = x[:, :, 3:4].astype(np.float32) / 255.0
        y = color * alpha + 255.0 * (1.0 - alpha)
        y = y.clip(0, 255).astype(np.uint8)
        return y

def process_image(
    image: Image.Image | str | None,
    mask: Image.Image | str | None,
    progress: gr.Progress = gr.Progress(),
) -> Image.Image | None:
    progress(0, desc="Preparing inputs...")
    if image is None or mask is None:
        return None
    
    if isinstance(mask, str):
        mask = Image.open(mask)
    if isinstance(image, str):
        image = Image.open(image)
    image = np.array(image)
    image = HWC3(image)

    result = simple_lama(image, mask)
    result.save("inpainted.png")
    return result

def resize_image(img: Image.Image, min_side_length: int = 768) -> Image.Image:
    if img.width <= min_side_length and img.height <= min_side_length:
        return img

    aspect_ratio = img.width / img.height
    if img.width < img.height:
        new_height = int(min_side_length / aspect_ratio)
        return img.resize((min_side_length, new_height))

    new_width = int(min_side_length * aspect_ratio)
    return img.resize((new_width, min_side_length))


async def process(
    image_and_mask: EditorValue | None,
    progress: gr.Progress = gr.Progress(),
) -> tuple[Image.Image, Image.Image] | None:
    if not image_and_mask:
        gr.Info("Please upload an image and draw a mask")
        return None


    image_np = image_and_mask["background"]
    image_np = cast(np.ndarray, image_np)

    if np.sum(image_np) == 0:
        gr.Info("Please upload an image")
        return None

    alpha_channel = image_and_mask["layers"][0]
    alpha_channel = cast(np.ndarray, alpha_channel)
    mask_np = np.where(alpha_channel[:, :, 3] == 0, 0, 255).astype(np.uint8)

    if np.sum(mask_np) == 0:
        gr.Info("Please mark the areas you want to remove")
        return None

    mask = Image.fromarray(mask_np)
    mask = resize_image(mask)

    image = Image.fromarray(image_np)
    image = resize_image(image)

    output = process_image(
        image,  
        mask,  
        progress,
    )

    if output is None:
        gr.Info("Processing failed")
        return None
    progress(100, desc="Processing completed")
    return image, output


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            image_and_mask = gr.ImageMask(
                label="Upload Image and Draw Mask",
                layers=False,
                show_fullscreen_button=False,
                sources=["upload"],
                show_download_button=False,
                interactive=True,
                height="full",
                width="full",
                brush=gr.Brush(default_size=75, colors=["#000000"], color_mode="fixed"),
                transforms=[],
            )


        with gr.Column():
            image_slider = ImageSlider(
                label="Result",
                interactive=False,
            )

            process_btn = gr.ClearButton(
                value="Run",
                variant="primary",
                size="lg",
                components=[image_slider],
            )

            process_btn.click(
                fn=lambda _: gr.update(interactive=False, value="Processing..."),
                inputs=[],
                outputs=[process_btn],
                api_name=False,
            ).then(
                fn=process,
                inputs=[
                    image_and_mask,
                ],
                outputs=[image_slider],
                api_name=False,
            ).then(
                fn=lambda _: gr.update(interactive=True, value="Run"),
                inputs=[],
                outputs=[process_btn],
                api_name=False,
            )


if __name__ == "__main__":
    demo.launch(
        debug=False,
        share=False,
        show_api=False,
    )