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

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 = []

RATIO_LABELS = {
    "1:1": "ratio_1",
    "9:16": "ratio_9_16",
}
RATIO_DISPLAY = list(RATIO_LABELS.keys())


def generate_image(
    prompt: str,
    style_id: str,
    input_image: Any,
    ratio_label: 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.")
    if ratio_label not in RATIO_LABELS:
        raise gr.Error("Please select a valid ratio.")

    ratio = RATIO_LABELS[ratio_label]

    image_path = None
    if input_image is not None:
        temp_path = "/tmp/somnium_input.png"
        if isinstance(input_image, Image.Image):
            input_image.save(temp_path)
        else:
            Image.fromarray(input_image).save(temp_path)
        image_path = temp_path
        logging.info(f"Edit mode: image saved to {temp_path}")

    try:
        logging.info(
            f"Generating image | style: '{style_id}' | ratio: '{ratio}' | edit: {image_path is not None}"
        )
        image = Somnium.Generate(
            prompt,
            STYLES_DICT[style_id],
            image_path,
            ratio=ratio
        )
        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 or edit your image.
        """
    )

    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[12] if STYLE_CHOICES else None
    )
    ratio_radio = gr.Radio(
        choices=RATIO_DISPLAY,
        label="Aspect Ratio:",
        value="9:16"
    )
    input_image = gr.Image(
        label="Input Image (optional — for editing):",
        type="pil",
        sources=["upload", "clipboard"]
    )
    generate_btn = gr.Button("Generate Image", variant="primary")
    image_output = gr.Image(
        label="Output",
        buttons=["download"]
    )

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

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