import gradio as gr
import time
from collections import defaultdict
from dotenv import load_dotenv
import os
from openai import OpenAI
import random

load_dotenv()
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

if client.api_key is None:
    raise ValueError("Die Umgebungsvariable OPENAI_API_KEY ist nicht gesetzt.")

APP_PASSWORD = os.getenv("APP_PASSWORD")
if APP_PASSWORD is None:
    raise ValueError("Die Umgebungsvariable APP_PASSWORD ist nicht gesetzt.")


GENERATION_TIMEOUT = float(os.getenv("GENERATION_TIMEOUT"))
if GENERATION_TIMEOUT is None:
    GENERATION_TIMEOUT = 60

last_generation_time = defaultdict(float)

def enhance_prompt(original_prompt):
    try:
        response = client.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "system", "content": "Verbessere den folgenden Bildgenerierungsprompt, um ihn detaillierter und kreativer zu machen."},
                {"role": "user", "content": original_prompt}
            ],
            max_tokens=100
        )
        enhanced_prompt = response.choices[0].message.content.strip()
        return enhanced_prompt
    except Exception as e:
        print("Ein Fehler ist beim Verbessern des Prompts aufgetreten:", e)
        return original_prompt

def generate_image(prompt, use_magic_prompt, style, user_id, password):
    global last_generation_time
    
    if password != APP_PASSWORD:
        return "Falsches Passwort. Bitte versuche es erneut.", None

    current_time = time.time()
    if current_time - last_generation_time[user_id] < GENERATION_TIMEOUT:
        return f"Bitte warte {GENERATION_TIMEOUT} Sekunden zwischen den Bildgenerierungen.", None

    if use_magic_prompt:
        prompt = enhance_prompt(prompt)
    
    print("Finaler Prompt:", prompt)

    try:
        response = client.images.generate(
            model="dall-e-3",
            prompt=prompt,
            size="1024x1024",
            quality="standard",
            style=style,
            n=1
        )
        
        image_url = response.data[0].url
        
        # Update the last generation time for this user
        last_generation_time[user_id] = current_time
        
        return prompt, image_url
    except Exception as e:
        error_message = str(e)
        if "content_policy_violation" in error_message:
            return "Leider können wir deinen Prompt nicht verarbeiten, da er evtl. gegen die Content Policy von OpenAI verstösst. Wenn du denkst, dass das ein Fehler ist, versuche den Prompt so umzuformulieren, dass er jugendfreundlich ist.", None
        print("Ein Fehler ist aufgetreten:", e)
        return f"Ein Fehler ist beim Generieren des Bildes aufgetreten: {error_message}", None

examples = [
    "Astronaut im Dschungel, kalte Farbpalette, gedämpfte Farben, detailliert, 8k",
    "Ein Astronaut reitet auf einem grünen Pferd",
    "Ein köstliches Ceviche-Käsekuchenstück",
    "Ein wunderschönes Glumanda in einer blühenden Wiese",
    "Ein futuristisches Stadtbild bei Nacht mit Neonlichtern",
    "Ein Roboter, der in einem antiken Tempel tanzt",
    "Ein magisches Einhorn, das über einen Regenbogen springt",
    "Ein mittelalterlicher Ritter, der gegen einen Drachen kämpft"
]

with gr.Blocks() as demo:
    gr.Markdown("# codora DALL-E 3 Bildgenerator")
    gr.Markdown(f"Gib einen Bildprompt ein und verwende optional die magische Prompt-Funktion, um ihn zu verbessern. Du kannst nur ein Bild alle {GENERATION_TIMEOUT} Sekunden generieren.")

    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(label="Bildprompt")
            use_magic_prompt = gr.Checkbox(label="Magischen Prompt verwenden")
            style = gr.Radio(
                label="Stil",
                choices=["vivid", "natural"],
                value="vivid"
            )
            user_id = gr.Textbox(label="Benutzer-ID", visible=False)
            password = gr.Textbox(label="App Passwort", type="password")
            generate_button = gr.Button("Bild generieren")

        with gr.Column(scale=1):
            final_prompt = gr.Textbox(label="Finaler Prompt")
            generated_image = gr.Image(label="Generiertes Bild", scale=0.5)  # Set scale to 0.5 for 50% width

    # Place examples below the generation button
    gr.Examples(
        examples=examples,
        inputs=[prompt]
    )

    generate_button.click(
        fn=generate_image,
        inputs=[prompt, use_magic_prompt, style, user_id, password],
        outputs=[final_prompt, generated_image]
    )

demo.launch()