import gradio as gr
from PIL import Image
import os
import time
import json
from fastapi import FastAPI
from pydantic import BaseModel

from src.utils import change_background, matte

app = FastAPI()

class ImageRequest(BaseModel):
    image: str  # Base64 encoded image
    background_color: str

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:
        background_color = "Transparent (PNG)"

    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)

@app.post("/process")
def api_process_images(request: ImageRequest):
    image_data = request.image
    background_color = request.background_color
    
    # Convert base64 to PIL image
    img_input = Image.open(BytesIO(base64.b64decode(image_data)))
    
    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_matte = matte(img_input)
    img_output = change_background(img_input, img_matte, background_alpha=alpha, background_hex=hexmap[background_color])
    
    buffered = BytesIO()
    img_output.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
    
    return {"image": img_str}

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", value="Transparent (PNG)")
    
    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()