import gradio as gr
from PIL import Image
import os
import time

from src.utils import change_background, matte

def process_images(uploaded_files, 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
    
    results = []
    times = []
    for uploaded_file in uploaded_files:
        start_time = time.time()
        img_input = Image.open(uploaded_file)  # Open each file
        img_matte = matte(img_input)
        img_output = change_background(img_input, img_matte, background_alpha=alpha, background_hex=hexmap[background_color])
        end_time = time.time()
        
        processing_time = end_time - start_time
        times.append(f"{os.path.basename(uploaded_file)}: {processing_time:.2f} seconds")
        results.append(img_output)
    
    return results, "\n".join(times)

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():
        folder_input = gr.Files(type="filepath", label="Upload your photos here")
        bg_color = gr.Dropdown(choices=["Transparent (PNG)", "White", "Black", "Green", "Red", "Blue"], label="Choose background color")
    
    output_images = gr.Gallery(label="Processed Images")
    processing_times = gr.Textbox(label="Processing Times")
    
    btn = gr.Button("Submit")
    
    btn.click(fn=process_images, inputs=[folder_input, bg_color], outputs=[output_images, processing_times])
    
    gr.Examples(
        examples=[["assets/demo.jpg", "Transparent (PNG)"]],
        inputs=[folder_input, bg_color]
    )

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