import gradio as gr
import numpy as np
import random
import torch
from diffusers import DiffusionPipeline, StableDiffusionImg2ImgPipeline
from PIL import Image

# -----------------------------
# CPU MODE ONLY
# -----------------------------
device = "cpu"
torch_dtype = torch.float32

MODEL_ID = "runwayml/stable-diffusion-v1-5"

# txt2img pipeline
txt2img_pipe = DiffusionPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch_dtype,
    low_cpu_mem_usage=True,
)
txt2img_pipe.to(device)

# img2img pipeline
img2img_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch_dtype,
    low_cpu_mem_usage=True,
)
img2img_pipe.to(device)

MAX_SEED = np.iinfo(np.int32).max


# -----------------------------
# Prompt builder
# -----------------------------
def build_prompt(prompt: str, style: str, category: str) -> str:
    style_map = {
        "Tanpa gaya": "",
        "Studio": "product photography, clean studio background, soft lighting, high quality",
        "E-commerce": "white background, catalog photo, sharp, high quality, tokopedia, shopee",
        "Pastel": "pastel colors, soft light, aesthetic instagram style",
        "Lifestyle": "realistic lifestyle photography, natural light",
        "Model Talent": "professional model, commercial photoshoot, studio lighting, natural pose, realistic skin texture, high quality",
    }

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

    s = style_map.get(style, "")
    c = category_map.get(category, "")

    parts = [prompt, s, c]
    return ", ".join([p for p in parts if p])


# -----------------------------
# Auto prompt generator
# -----------------------------
def auto_prompt(category: str) -> str:
    templates = {
        "Skincare": "Serum skincare botol kaca premium, tampilan mewah, cocok untuk iklan Instagram",
        "Makanan/Minuman": "Minuman energi rasa lemon, efek splash, gaya promosi e-commerce",
        "Fashion": "Sepatu running sport, tampilan katalog, background putih bersih",
        "Elektronik": "Headphone wireless modern, lighting studio, tampilan premium",
        "Umum": "Produk premium dengan lighting studio dan background bersih",
    }
    return templates.get(category, "Produk premium dengan lighting studio dan background bersih")


# -----------------------------
# Inference
# -----------------------------
def generate(
    mode,
    prompt,
    negative_prompt,
    seed,
    randomize_seed,
    width,
    height,
    guidance_scale,
    steps,
    style,
    category,
    num_images,
    init_image,
    strength,
):
    if not prompt:
        raise gr.Error("Prompt tidak boleh kosong.")

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator(device=device).manual_seed(seed)
    full_prompt = build_prompt(prompt, style, category)

    images = []

    if mode == "Text to Image":
        for _ in range(num_images):
            out = txt2img_pipe(
                prompt=full_prompt,
                negative_prompt=negative_prompt or None,
                width=width,
                height=height,
                guidance_scale=guidance_scale,
                num_inference_steps=steps,
                generator=generator,
            )
            images.append(out.images[0])
    else:  # Image to Image
        if init_image is None:
            raise gr.Error("Upload gambar produk terlebih dahulu.")

        init_image = init_image.convert("RGB").resize((width, height))

        for _ in range(num_images):
            out = img2img_pipe(
                prompt=full_prompt,
                negative_prompt=negative_prompt or None,
                image=init_image,
                strength=strength,
                guidance_scale=guidance_scale,
                num_inference_steps=steps,
                generator=generator,
            )
            images.append(out.images[0])

    return images, seed


# -----------------------------
# UI
# -----------------------------
with gr.Blocks(title="RuangAI – Product Visualizer CPU") as demo:
    gr.Markdown(
        """
        # 🧴 RuangAI – Product Visualizer (CPU Mode)

        - **Text to Image**: buat visual produk dari deskripsi  
        - **Image to Image**: upload foto produk lalu buat versi promosi  
        - Pilih **gaya visual** dan **kategori produk**  
        - Gaya **Model Talent** akan menambahkan visualisasi seorang model di hasil gambar
        """
    )

    mode = gr.Radio(
        ["Text to Image", "Image to Image"],
        value="Text to Image",
        label="Mode",
    )

    with gr.Row():
        category = gr.Dropdown(
            ["Umum", "Skincare", "Makanan/Minuman", "Fashion", "Elektronik"],
            value="Umum",
            label="Kategori Produk",
        )
        auto_btn = gr.Button("Auto Prompt ✨")

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

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

    init_image = gr.Image(
        label="Upload Gambar (untuk Image to Image)",
        type="pil",
    )

    with gr.Row():
        style = gr.Dropdown(
            ["Tanpa gaya", "Studio", "E-commerce", "Pastel", "Lifestyle", "Model Talent"],
            value="Studio",
            label="Gaya visual",
        )
        num_images = gr.Slider(
            1, 4, value=1, step=1, label="Jumlah gambar"
        )

    gallery = gr.Gallery(
        label="Hasil",
        columns=2,
        height=512,
    )

    with gr.Accordion("Advanced Settings", open=False):
        negative_prompt = gr.Textbox(
            label="Negative prompt",
            placeholder="Contoh: blur, low quality, watermark, text, logo",
        )
        seed = gr.Slider(
            0, MAX_SEED, value=0, step=1, label="Seed"
        )
        randomize_seed = gr.Checkbox(
            True, label="Randomize seed"
        )
        width = gr.Slider(
            256, 768, value=512, step=32, label="Width"
        )
        height = gr.Slider(
            256, 768, value=512, step=32, label="Height"
        )
        guidance_scale = gr.Slider(
            0, 10, value=7, step=0.5, label="Guidance"
        )
        steps = gr.Slider(
            5, 40, value=25, step=1, label="Steps"
        )
        strength = gr.Slider(
            0.1, 1.0, value=0.6, step=0.05, label="Strength (img2img)"
        )

    run_btn = gr.Button("Generate 🚀")

    run_btn.click(
        generate,
        inputs=[
            mode,
            prompt,
            negative_prompt,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            steps,
            style,
            category,
            num_images,
            init_image,
            strength,
        ],
        outputs=[gallery, seed],
    )

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