import requests
import gradio as gr
from PIL import Image
import io
import os

# API settings
API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-3-medium-diffusers"
headers = {"Authorization": f"Bearer {os.getenv('HF_API_KEY')}"}

def query(payload):
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.content

def generate_room_plan(space_type, prompt):
    # Validate the prompt based on selected space type
    valid_prompt = False
    if space_type == "office" and "office" in prompt:
        valid_prompt = True
    elif space_type == "house" and "house" in prompt:
        valid_prompt = True
    elif space_type == "room" and "room" in prompt:
        valid_prompt = True
    
    if not valid_prompt:
        return "Error: The prompt does not match the selected space type or is invalid."

    image_bytes = query({"inputs": prompt})
    
    # Convert the byte content to an image
    image = Image.open(io.BytesIO(image_bytes))
    
    return image

with gr.Blocks() as app:
    gr.Markdown("# Space Plan Generator")

    with gr.Row():
        space_type = gr.Dropdown(
            choices=["office", "house", "room"],
            label="Select Space Type"
        )
        user_prompt = gr.Textbox(label="Enter the dimensions and other details")

    output = gr.Image(label="Generated Space Plan")
    error_message = gr.Markdown(label="")

    def on_button_click(space_type, prompt):
        result = generate_room_plan(space_type, prompt)
        if isinstance(result, str):  # If the result is an error message
            error_message.update(result)
            return None
        else:
            error_message.update("")
            return result

    gr.Button("Generate Space Plan").click(
        on_button_click,
        inputs=[space_type, user_prompt],
        outputs=[output]
    )

app.launch()