import logging
from typing import Any
import gradio as gr
from somnium import Somnium

logging.basicConfig(
    level=logging.INFO, 
    format="%(asctime)s - %(levelname)s - %(message)s"
)

try:
    STYLES_DICT = Somnium.Styles()
    STYLE_CHOICES = list(STYLES_DICT.keys())
except Exception as e:
    logging.error(f"Failed to load Somnium styles: {e}")
    STYLES_DICT = {}
    STYLE_CHOICES = []

def generate_image(prompt: str, style_id: str) -> Any:
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a valid prompt.")
    if not style_id or style_id not in STYLES_DICT:
        raise gr.Error("Please select a valid style.")

    try:
        logging.info(f"Generating image with style: '{style_id}'")
        image = Somnium.Generate(prompt, STYLES_DICT[style_id])
        return image
    except Exception as e:
        logging.error(f"Generation failed: {e}", exc_info=True)
        raise gr.Error("Generation failed. Please try a different prompt or check for NSFW content.")


with gr.Blocks(title="Somnium Image Generator") as demo:
    
    gr.Markdown(
        """
        # 🌌 Somnium Image Generator
        Enter a prompt and select a style to generate your image.
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            prompt_input = gr.Textbox(
                label="Enter Prompt:", 
                placeholder="A futuristic city at sunset...",
                lines=5
            )
            style_dropdown = gr.Dropdown(
                choices=STYLE_CHOICES, 
                label="Select Style:",
                value=STYLE_CHOICES[0] if STYLE_CHOICES else None
            )
            generate_btn = gr.Button("Generate Image", variant="primary")
            
        with gr.Column(scale=1):
            image_output = gr.Image(
                label="Output", 
                buttons=["download"] 
            )

    generate_btn.click(
        fn=generate_image,
        inputs=[prompt_input, style_dropdown],
        outputs=image_output
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft())