import gradio as gr
from PIL import Image
from io import BytesIO

from src.utils import change_background, matte

def process_image(uploaded_file, background_color):
    # Add debugging statement
    print(f"background_color received: {background_color}")

    # Default value check
    if background_color is None:
        return "Please select a background color."

    hexmap = {
        "Transparent (PNG)": "#000000",
        "Black": "#000000",
        "White": "#FFFFFF",
        "Green": "#22EE22",
        "Red": "#EE2222",
        "Blue": "#2222EE",
    }
    alpha = 0.0 if background_color == "Transparent (PNG)" else 1.0

    img_input = uploaded_file  # No need to open, it's already a PIL image
    img_matte = matte(img_input)
    img_output = change_background(img_input, img_matte, background_alpha=alpha, background_hex=hexmap[background_color])
    
    return img_output

with gr.Blocks() as demo:
    gr.Markdown("""
    # AI Photo Background Removal
    
    You want to remove your photo background, but don't have the time and effort to learn photo editing skills?
    **This app will change or remove your photo background, in seconds.**
    """)
    
    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload your photo here")
        bg_color = gr.Dropdown(choices=["Transparent (PNG)", "White", "Black", "Green", "Red", "Blue"], label="Choose background color")
    
    output_image = gr.Image(type="pil", label="Processed Image")
    
    btn = gr.Button("Submit")
    
    btn.click(fn=process_image, inputs=[image_input, bg_color], outputs=output_image)
    
    gr.Examples(
        examples=[["assets/demo.jpg", "Transparent (PNG)"]],
        inputs=[image_input, bg_color]
    )

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