"""
BatikAI — Gradio Web Application
Interactive batik motif recognition with cultural insights.
"""

import gradio as gr
from PIL import Image
import os

# ─────────────────────────────────────────────
# MODEL LOADING (lazy, done once at startup)
# ─────────────────────────────────────────────
MODEL_DIR = "./model"
_model_state = {}

def get_model():
    if not _model_state:
        from inference import load_model
        model, device, transform, class_names, config = load_model(MODEL_DIR)
        _model_state["model"]       = model
        _model_state["device"]      = device
        _model_state["transform"]   = transform
        _model_state["class_names"] = class_names
        _model_state["config"]      = config
    return (
        _model_state["model"],
        _model_state["device"],
        _model_state["transform"],
        _model_state["class_names"],
        _model_state["config"],
    )


# ─────────────────────────────────────────────
# PREDICTION FUNCTION
# ─────────────────────────────────────────────
def predict_batik(image: Image.Image):
    """Run inference and return structured Gradio outputs."""
    if image is None:
        return {}, "", ""

    from inference import predict

    model, device, transform, class_names, config = get_model()
    predictions, info = predict(image, model, device, transform, class_names, top_k=5)

    # Build confidence dict for gr.Label
    label_confidences = {name: conf for name, conf in predictions}

    # Cultural info
    top_class   = predictions[0][0].replace("_", " ").title()
    confidence  = predictions[0][1]
    origin      = info.get("origin",  "—")
    meaning     = info.get("meaning", "—")
    usage       = info.get("usage",   "—")
    fun_fact    = info.get("fun_fact","—")

    header = f"## 🎨 {top_class}  —  {confidence*100:.1f}% confident"

    cultural_md = f"""
### 📍 Asal Daerah
{origin}

### 💬 Makna & Filosofi
{meaning}

### 👘 Penggunaan Tradisional
{usage}

### ✨ Fun Fact
{fun_fact}
"""

    return label_confidences, header, cultural_md


# ─────────────────────────────────────────────
# GRADIO UI
# ─────────────────────────────────────────────
TITLE       = "🏵️ BatikAI — Pengenalan Motif Batik Indonesia"
DESCRIPTION = """
Upload foto kain batik dan BatikAI akan langsung mengenali motifnya beserta latar belakang budayanya!

Model dilatih dengan **ResNet-50** yang di-fine-tune pada dataset [Batik Indonesia](https://huggingface.co/datasets/muhammadsalmanalfaridzi/Batik-Indonesia)
mencakup **10 motif batik** autentik:  
*Parang, Kawung, Mega Mendung, Truntum, Sekar Jagad, Ceplok, Sidomukti, Lereng, Nitik, Sogan*
"""

THEME = gr.themes.Soft(
    primary_hue="amber",
    secondary_hue="orange",
    neutral_hue="stone",
    font=gr.themes.GoogleFont("Inter"),
)

with gr.Blocks(theme=THEME, title="BatikAI") as demo:
    gr.Markdown(f"# {TITLE}")
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        # ── Left column: input
        with gr.Column(scale=1):
            image_input = gr.Image(
                type="pil",
                label="📸 Upload Gambar Batik",
                height=340,
            )
            btn = gr.Button("🔍 Kenali Motif", variant="primary", size="lg")

        # ── Right column: outputs
        with gr.Column(scale=1):
            result_header = gr.Markdown("*Hasil prediksi akan muncul di sini...*")
            label_output  = gr.Label(
                num_top_classes=5,
                label="Top-5 Prediksi",
            )

    # Cultural info panel
    with gr.Row():
        cultural_output = gr.Markdown(label="Informasi Budaya")

    # ── Examples (only shown if the examples folder exists)
    examples_dir = "./examples"
    if os.path.isdir(examples_dir):
        example_files = [
            [os.path.join(examples_dir, f)]
            for f in os.listdir(examples_dir)
            if f.lower().endswith((".jpg", ".jpeg", ".png", ".webp"))
        ]
        if example_files:
            gr.Examples(
                examples=example_files,
                inputs=image_input,
                label="Contoh Gambar Batik",
            )

    # ── Wire up
    btn.click(
        fn=predict_batik,
        inputs=[image_input],
        outputs=[label_output, result_header, cultural_output],
    )
    image_input.change(
        fn=predict_batik,
        inputs=[image_input],
        outputs=[label_output, result_header, cultural_output],
    )

    gr.Markdown(
        """
---
<center>
Made with ❤️ for Indonesian cultural heritage · Model: <b>ResNet-50</b> · Dataset: HuggingFace Batik-Indonesia
</center>
"""
    )


if __name__ == "__main__":
    demo.launch(
        share=False,
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
    )