import numpy as np
import gradio as gr
from PIL import Image
from io import BytesIO
from copy import deepcopy
from src.core import process_inpaint
from huggingface_hub import login
import os
login(os.getenv("HF_TOKEN"))

def process_image(img_input, mask_input, brush_size):
    img_input = Image.open(BytesIO(img_input)).convert("RGBA")
    mask_input = Image.open(BytesIO(mask_input)).convert("RGBA")
    
    max_size = 2000
    img_width, img_height = img_input.size
    if img_width > max_size or img_height > max_size:
        if img_width > img_height:
            new_width = max_size
            new_height = int((max_size / img_width) * img_height)
        else:
            new_height = max_size
            new_width = int((max_size / img_height) * img_width)
        img_input = img_input.resize((new_width, new_height))
    
    im = np.array(mask_input.resize(img_input.size))
    background = np.where(
        (im[:, :, 0] == 0) & 
        (im[:, :, 1] == 0) & 
        (im[:, :, 2] == 0)
    )
    drawing = np.where(
        (im[:, :, 0] == 255) & 
        (im[:, :, 1] == 0) & 
        (im[:, :, 2] == 255)
    )
    im[background] = [0, 0, 0, 255]
    im[drawing] = [0, 0, 0, 0]
    
    output = process_inpaint(np.array(img_input), np.array(im))
    img_output = Image.fromarray(output).convert("RGB")
    
    output_buffer = BytesIO()
    img_output.save(output_buffer, format="PNG")
    return output_buffer.getvalue()

demo = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(type="bytes", label="Upload Image"),
        gr.Image(type="bytes", tool="sketch", label="Draw Mask"),
        gr.Slider(1, 100, value=50, label="Brush Size")
    ],
    outputs=gr.Image(type="file", label="Output Image"),
    title="Object Remover",
)

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