import os
from io import BytesIO

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

# =========================
# DEVICE
# =========================

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32

print(f"[INFO] Using device: {DEVICE}, dtype: {DTYPE}")


# =========================
# MODEL CONFIG
# =========================

MODEL_OPTIONS = {
    "Realistic Vision v5.1": "SG161222/Realistic_Vision_V5.1_noVAE",
    "Stable Diffusion 1.5": "runwayml/stable-diffusion-v1-5",
    "DreamShaper 8": "Lykon/dreamshaper-8",
}

PIPELINES = {}


def get_pipeline(model_name: str) -> StableDiffusionPipeline:
    if model_name in PIPELINES:
        return PIPELINES[model_name]

    repo_id = MODEL_OPTIONS[model_name]
    print(f"[INFO] Loading model: {model_name} ({repo_id})")

    pipe = StableDiffusionPipeline.from_pretrained(
        repo_id,
        torch_dtype=DTYPE,
        safety_checker=None,
    )

    pipe = pipe.to(DEVICE)

    if DEVICE == "cuda":
        pipe.enable_xformers_memory_efficient_attention()

    PIPELINES[model_name] = pipe
    return pipe


# =========================
# PROMPT SYSTEM
# =========================

def auto_prompt(category: str) -> str:
    templates = {
        "Skincare": "Serum skincare botol kaca premium, lighting studio, aesthetic clean look",
        "Makanan/Minuman": "Minuman segar dengan efek splash, lighting vibrant, cocok untuk iklan",
        "Fashion": "Pakaian atau sepatu fashion modern, lighting studio, katalog e-commerce",
        "Elektronik": "Headphone wireless premium, lighting studio, tampilan high-end",
        "Umum": "Produk premium dengan lighting studio dan background bersih",
    }
    return templates.get(category, templates["Umum"])


def build_prompt(prompt: str, style: str, category: str, with_model: bool) -> str:
    style_map = {
        "Tanpa gaya": "",
        "Studio": "studio lighting, clean background, high quality product photography",
        "E-commerce": "white background, catalog photo, sharp, high quality",
        "Pastel": "pastel colors, soft light, aesthetic instagram style",
        "Lifestyle": "realistic lifestyle photography, natural light",
    }

    category_map = {
        "Umum": "",
        "Skincare": "skincare product, glossy bottle, beauty aesthetic",
        "Makanan/Minuman": "food photography, appetizing, vibrant lighting",
        "Fashion": "fashion product, textile detail, clean lighting",
        "Elektronik": "electronic product, reflective surface, studio lighting",
    }

    model_snippet = (
        "professional model, commercial photoshoot, natural pose, holding the product"
        if with_model else ""
    )

    parts = [
        prompt,
        style_map.get(style, ""),
        category_map.get(category, ""),
        model_snippet,
        "high quality, 4k, detailed",
    ]

    return ", ".join([p for p in parts if p])


# =========================
# GENERATION
# =========================

def run(prompt, category, style, with_model, model_choice, steps, guidance, seed):
    if not prompt or prompt.strip() == "":
        prompt = auto_prompt(category)

    full_prompt = build_prompt(prompt, style, category, with_model)

    pipe = get_pipeline(model_choice)

    generator = None
    if seed is not None and seed != "":
        try:
            seed_int = int(seed)
            generator = torch.Generator(device=DEVICE).manual_seed(seed_int)
        except ValueError:
            generator = None

    result = pipe(
        full_prompt,
        num_inference_steps=int(steps),
        guidance_scale=float(guidance),
        generator=generator,
    )

    img: Image.Image = result.images[0]
    return img


# =========================
# GRADIO UI
# =========================

with gr.Blocks(title="RuangAI – Product Visualizer (Diffusers)") as demo:
    gr.Markdown("""
    # 🧴 RuangAI – Product Visualizer (Level 2 – Diffusers Lokal)
    Tiga model lokal: Realistic Vision v5.1, Stable Diffusion 1.5, DreamShaper 8  
    **Catatan:** di CPU akan agak lambat, sabar sebentar saat generate 🙏
    """)

    with gr.Row():
        with gr.Column():
            model_choice = gr.Dropdown(
                list(MODEL_OPTIONS.keys()),
                value="Realistic Vision v5.1",
                label="Pilih Model",
            )

            category = gr.Dropdown(
                ["Umum", "Skincare", "Makanan/Minuman", "Fashion", "Elektronik"],
                value="Umum",
                label="Kategori Produk",
            )

            style = gr.Dropdown(
                ["Tanpa gaya", "Studio", "E-commerce", "Pastel", "Lifestyle"],
                value="Studio",
                label="Gaya Visual",
            )

            with_model = gr.Checkbox(
                label="Tambahkan Model Talent (Manusia)",
                value=False,
            )

            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Deskripsi produk / ide visual...",
                lines=3,
            )

            with gr.Row():
                auto_btn = gr.Button("Auto Prompt ✨")
                generate_btn = gr.Button("Generate 🚀")

            steps = gr.Slider(
                minimum=10,
                maximum=40,
                value=25,
                step=1,
                label="Inference Steps",
            )

            guidance = gr.Slider(
                minimum=3.0,
                maximum=12.0,
                value=7.5,
                step=0.5,
                label="Guidance Scale",
            )

            seed = gr.Textbox(
                label="Seed (opsional, untuk hasil konsisten)",
                placeholder="Kosongkan untuk random",
            )

        with gr.Column():
            output_image = gr.Image(label="Hasil", type="pil")

    auto_btn.click(auto_prompt, inputs=[category], outputs=[prompt])

    generate_btn.click(
        run,
        inputs=[prompt, category, style, with_model, model_choice, steps, guidance, seed],
        outputs=[output_image],
    )

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