import os
import requests
from io import BytesIO
from PIL import Image
import gradio as gr
from dotenv import load_dotenv

# =========================
# LOAD ENV
# =========================

load_dotenv()

HF_TOKEN = os.environ.get("HF_TOKEN")
HF_MODEL = os.environ.get("HF_MODEL", "playgroundai/playground-v2.5")

API_URL = "https://api-inference.huggingface.co/models/stabilityai/sdxl-turbo"

if not HF_TOKEN:
    print("[WARNING] HF_TOKEN belum di-set di .env")


# =========================
# 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": "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,
    ]

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


# =========================
# HUGGINGFACE API CALL
# =========================

def call_huggingface(prompt: str):
    headers = {
        "Authorization": f"Bearer {HF_TOKEN}",
        "Content-Type": "application/json",
    }

    payload = {
        "inputs": prompt
    }

    response = requests.post(API_URL, headers=headers, json=payload)
    response.raise_for_status()

    img_bytes = response.content
    img = Image.open(BytesIO(img_bytes))

    return img


# =========================
# MAIN GENERATION
# =========================

def run(prompt, category, style, with_model):
    full_prompt = build_prompt(prompt, style, category, with_model)
    img = call_huggingface(full_prompt)
    return img


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

with gr.Blocks(title="RuangAI – Product Visualizer (Level 2)") as demo:
    gr.Markdown("""
    # 🧴 RuangAI – Product Visualizer (Level 2)
    Playground v2.5 (HuggingFace Inference API)
    """)

    with gr.Row():
        with gr.Column():
            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,
            )

            auto_btn = gr.Button("Auto Prompt ✨")
            auto_btn.click(auto_prompt, inputs=[category], outputs=[prompt])

            generate_btn = gr.Button("Generate 🚀")

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

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

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