import gradio as gr
import fal_client
from dotenv import load_dotenv
import requests

load_dotenv()

def process_image(image, prompt):
    image_url = fal_client.upload_file(image)
    arguments = {
        "image_url": image_url,
        "prompt": prompt
    }
    handler = fal_client.submit(
        "workflows/acemetrics/bgrep",
        arguments=arguments
    )
    result = fal_client.result("workflows/acemetrics/bgrep", handler.request_id)
    result_url = result['images'][0]['url']
    response = requests.get(result_url)
    temp_output = "temp_output.png"
    with open(temp_output, "wb") as f:
        f.write(response.content)
    
    return temp_output

with gr.Blocks(title="AI Background Replacement") as demo:
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="Upload Image", type="filepath")
            prompt_input = gr.Textbox(label="Background Prompt", lines=2)
            process_btn = gr.Button("Replace Background", variant="primary")
        
        with gr.Column():
            output_image = gr.Image(label="Result")
    
    gr.Examples(
        examples=[
            ["examples/person_hills.jpg", "Snow-covered Alps mountains with falling snow"],
            ["examples/person_wall.jpg", "Dense forest with trees and natural lighting"],
            ["examples/person_street.jpg", "Serene ocean beach with soft waves"],
        ],
        inputs=[image_input, prompt_input],
        fn=process_image,
        outputs=output_image,
    )
    
    process_btn.click(
        fn=process_image,
        inputs=[image_input, prompt_input],
        outputs=output_image
    )

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