"""
FHDR Uncensored Text-to-Image Generator
Using kpsss34/FHDR_Uncensored via Diffusers (FluxPipeline)
"""

import logging
from typing import Optional

import gradio as gr
from PIL import Image
import torch
from diffusers import FluxPipeline

# ── Logging ────────────────────────────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger(__name__)

# ── Model Loader ───────────────────────────────────────────────────────────────
pipe = None

def load_model():
    global pipe
    if pipe is None:
        logger.info("Loading FHDR model...")
        pipe = FluxPipeline.from_pretrained(
            "kpsss34/FHDR_Uncensored",
            torch_dtype=torch.bfloat16
        )
        pipe.enable_model_cpu_offload()
        logger.info("Model loaded successfully!")

# ── Styles ─────────────────────────────────────────────────────────────────────
STYLES: dict[str, str] = {
    "None": "",
    "Realistic": "photorealistic, ultra-detailed, 8K, sharp focus",
    "Anime": "anime style, vibrant colors, detailed linework",
    "Cyberpunk": "cyberpunk, neon lights, futuristic city",
    "Fantasy": "epic fantasy, cinematic lighting, concept art",
}

EXAMPLE_PROMPTS = [
    ["A dragon flying over mountains", "Fantasy", 4.0, 768, 768],
    ["Cyberpunk city at night", "Cyberpunk", 4.0, 768, 768],
]

# ── Helper ─────────────────────────────────────────────────────────────────────
def apply_style_to_prompt(user_prompt: str, style_name: str) -> str:
    style_keywords = STYLES.get(style_name, "")
    return f"{user_prompt}, {style_keywords}" if style_keywords else user_prompt

# ── Core Generation ────────────────────────────────────────────────────────────
def generate_image(
    user_prompt: str,
    style_name: str,
    guidance_scale: float,
    width: int,
    height: int,
    progress: gr.Progress = gr.Progress(),
) -> tuple[Optional[Image.Image], str]:

    if not user_prompt.strip():
        return None, "⚠️ Please enter a prompt."

    try:
        progress(0.2, desc="Loading model...")
        load_model()

        final_prompt = apply_style_to_prompt(user_prompt, style_name)

        progress(0.5, desc="Generating image...")

        image = pipe(
            final_prompt,
            height=height,
            width=width,
            guidance_scale=guidance_scale,
            num_inference_steps=30,
            max_sequence_length=512,
            generator=torch.Generator("cpu").manual_seed(0)
        ).images[0]

        progress(1.0, desc="Done!")
        return image, "✅ Image generated successfully!"

    except Exception as e:
        logger.error("Error: %s", e, exc_info=True)
        return None, f"❌ Error: {str(e)}"

# ── UI ─────────────────────────────────────────────────────────────────────────
def build_interface() -> gr.Blocks:
    with gr.Blocks(title="FHDR Image Generator") as demo:

        gr.HTML("""
        <div style="text-align:center;">
            <h1>⚡ FHDR Uncensored Image Generator</h1>
            <p>Powered by kpsss34/FHDR_Uncensored (Flux)</p>
        </div>
        """)

        with gr.Row():
            with gr.Column():
                prompt_input = gr.Textbox(label="Prompt", lines=3)
                style_input = gr.Dropdown(
                    choices=list(STYLES.keys()),
                    value="None",
                    label="Style"
                )

                guidance_input = gr.Slider(1.0, 10.0, value=4.0, label="Guidance")

                with gr.Row():
                    width_input = gr.Slider(512, 1024, value=768, step=64, label="Width")
                    height_input = gr.Slider(512, 1024, value=768, step=64, label="Height")

                generate_btn = gr.Button("Generate")
                status_output = gr.Markdown()

            with gr.Column():
                image_output = gr.Image(label="Output")

        generate_btn.click(
            fn=generate_image,
            inputs=[prompt_input, style_input, guidance_input, width_input, height_input],
            outputs=[image_output, status_output],
        )

        gr.Examples(
            examples=EXAMPLE_PROMPTS,
            inputs=[prompt_input, style_input, guidance_input, width_input, height_input],
            outputs=[image_output, status_output],
            fn=generate_image,
        )

    return demo

# ── Run ────────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    demo = build_interface()
    demo.launch()