import os
import io
import base64
from PIL import Image
import gradio as gr
from openai import OpenAI

# Initialize the OpenAI client
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
client = OpenAI(api_key=OPENAI_API_KEY)

def inference(input_image: Image.Image) -> Image.Image:
    """
    Enhances a handwritten sketch into a clean, structured floor plan using OpenAI's image model.
    """
    try:
        # Convert the PIL image to a PNG byte buffer
        img_byte_arr = io.BytesIO()
        input_image.save(img_byte_arr, format="PNG")
        img_byte_arr.seek(0)
        img_byte_arr.name = "input.png"  # ✅ Give it a valid filename

        prompt = """
        Enhance this handwritten sketch into a clean, structured, computer-drawn floor plan design.
        """

        # Use the in-memory file directly
        result = client.images.edit(
            model="gpt-image-1",
            image=[img_byte_arr],
            prompt=prompt
        )

        # Decode the returned base64 image
        image_base64 = result.data[0].b64_json
        image_bytes = base64.b64decode(image_base64)

        # Convert bytes back to PIL Image
        enhanced_image = Image.open(io.BytesIO(image_bytes))

        return enhanced_image

    except Exception as e:
        print(f"Error during image enhancement: {e}")
        return input_image  # fallback to original image

# Gradio interface
iface = gr.Interface(
    fn=inference,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="Sketch to Floor Plan Pix2Pix",
    description="Upload a sketch image and the model will enhance it into a structured floor plan."
)

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