import gradio as gr
from openai import OpenAI
from PIL import Image
import requests
from io import BytesIO
import logging
import re
import os

logging.basicConfig(level=logging.INFO)


def sanitize_filename(prompt):
    sanitized = re.sub(r"[^\w\-]", "_", prompt)
    sanitized = re.sub(r"_+", "_", sanitized)
    sanitized = sanitized.strip("_")
    return sanitized[:50]


def generate_image(api_key, prompt, file_format, image_size):
    if not api_key:
        raise gr.Error("API key is required. Please enter your OpenAI API key.")

    if not prompt or prompt.strip() == "":
        raise gr.Error(
            "Prompt cannot be empty. Please enter a description for the image you want to generate."
        )

    client = OpenAI(api_key=api_key)

    try:
        logging.info(f"Attempting to generate image with prompt: {prompt}")
        response = client.images.generate(
            model="dall-e-3",
            prompt=prompt,
            size=image_size,
            quality="standard",
            n=1,
        )

        logging.info("API call successful")
        image_url = response.data[0].url
        logging.info(f"Image URL received: {image_url}")

        image_response = requests.get(image_url)
        img = Image.open(BytesIO(image_response.content))

        filename = f"{sanitize_filename(prompt)}.{file_format}"
        img_path = os.path.join("output", filename)
        os.makedirs("output", exist_ok=True)
        img.save(img_path, format=file_format.upper())

        logging.info(f"Image successfully generated and saved as {img_path}")
        return img_path
    except Exception as e:
        logging.error(f"Error occurred: {str(e)}")
        if "invalid_api_key" in str(e):
            raise gr.Error(
                "Invalid API key. Please check your OpenAI API key and try again."
            )
        else:
            raise gr.Error(f"An error occurred: {str(e)}")


with gr.Blocks() as demo:
    gr.Markdown("""
    # DALL-E 3 Image Generation

    Generate images using OpenAI's DALL-E 3 model.

    **Important:** You need an OpenAI API key to use this application. You can obtain one from [https://platform.openai.com/api-keys](https://platform.openai.com/api-keys).

    **Warning:** Your API key is sensitive information. This application does not store your key, but please handle it with care and do not share it with others.

    **Please Note:**
    1. Image generation typically takes around 30 seconds per image.
    2. Pricing varies based on image size. For the most up-to-date pricing information, please visit the [OpenAI Pricing Page](https://openai.com/api/pricing/).
    3. We are not responsible for any failures in image generation. Use this application at your own risk.
    """)

    with gr.Row():
        with gr.Column(scale=1):
            api_key_input = gr.Textbox(label="OpenAI API Key", type="password")
            prompt_input = gr.Textbox(
                label="Enter your prompt", placeholder="e.g., a white siamese cat"
            )
            file_format = gr.Radio(["webp", "png"], label="File Format", value="webp")
            image_size = gr.Radio(
                ["1024x1024", "1024x1792", "1792x1024"],
                label="Image Size",
                value="1024x1024",
            )
            submit_button = gr.Button("Generate Image")

        with gr.Column(scale=1):
            image_output = gr.Image(label="Generated Image")

    submit_button.click(
        generate_image,
        inputs=[api_key_input, prompt_input, file_format, image_size],
        outputs=image_output,
    )

    gr.Examples(
        examples=[
            "A white siamese cat",
            "A futuristic cityscape at night",
            "A serene mountain lake at sunrise",
            "An abstract painting inspired by jazz music",
        ],
        inputs=prompt_input,
    )

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