import os
import requests
import gradio as gr

FAL_KEY = os.environ.get("FAL_KEY")

FAL_MODELS = {
    "Flux Schnell (Gratis)": "fal-ai/flux-schnell",
    "Flux Dev (Lebih Detail)": "fal-ai/flux-dev",
}

def auto_prompt(category):
    templates = {
        "Skincare": "premium skincare bottle, studio lighting, glossy, aesthetic clean look",
        "Makanan/Minuman": "fresh drink with splash effect, vibrant lighting, commercial photo",
        "Fashion": "modern fashion product, catalog lighting, clean background",
        "Elektronik": "premium electronic product, reflective surface, studio lighting",
        "Umum": "premium product, studio lighting, clean background",
    }
    return templates.get(category, templates["Umum"])

def build_prompt(prompt, style, category, with_model):
    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])

def generate(prompt, category, style, with_model, model_choice):
    if not prompt:
        prompt = auto_prompt(category)

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

    url = f"https://fal.run/v1/{FAL_MODELS[model_choice]}"

    headers = {
        "Authorization": f"Key {FAL_KEY}",
        "Content-Type": "application/json",
    }

    payload = {
        "prompt": full_prompt,
        "image_size": "1024x1024"
    }

    # Step 1: Submit request
    response = requests.post(url, json=payload, headers=headers)
    data = response.json()

    # Case A: Error message
    if "error" in data:
        return f"Fal.ai Error: {data['error']}"

    # Case B: Direct image output
    if "images" in data:
        return data["images"][0]["url"]

    # Case C: Async mode → need to poll
    if "request_id" in data:
        request_id = data["request_id"]

        # Polling endpoint
        poll_url = f"https://fal.run/v1/requests/{request_id}"

        while True:
            poll = requests.get(poll_url, headers=headers).json()

            if poll.get("status") == "completed":
                return poll["images"][0]["url"]

            if poll.get("status") == "failed":
                return f"Fal.ai failed: {poll}"

            # Wait a bit
            import time
            time.sleep(1)

    # Unknown response
    return f"Unexpected response: {data}"

with gr.Blocks(title="RuangAI – Product Visualizer (Fal.ai)") as demo:
    gr.Markdown("# 🧴 RuangAI – Product Visualizer (Fal.ai Version)")

    with gr.Row():
        with gr.Column():
            model_choice = gr.Dropdown(
                list(FAL_MODELS.keys()),
                value="Flux Schnell (Gratis)",
                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...",
                lines=3
            )

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

        with gr.Column():
            output = gr.Image(label="Hasil")

    auto_btn.click(auto_prompt, inputs=[category], outputs=[prompt])
    generate_btn.click(
        generate,
        inputs=[prompt, category, style, with_model, model_choice],
        outputs=[output]
    )

demo.launch()