# app.py
# HF Spaces (Gradio 6.x) + TensorFlow Keras .keras image classifier
# UPDATE FINAL: Panduan tindakan hanya 1 (berdasarkan prediksi utama & threshold)

import os
import json
from pathlib import Path

import numpy as np
from PIL import Image
import gradio as gr
import tensorflow as tf

# -----------------------
# Config
# -----------------------
MODEL_PATH = os.environ.get("MODEL_PATH", "export/best_model.keras")
CLASSES_PATH = os.environ.get("CLASSES_PATH", "export/class_names.json")
IMG_SIZE = int(os.environ.get("IMG_SIZE", "224"))
TOP_K = int(os.environ.get("TOP_K", "3"))

# Ambang minimum agar panduan spesifik muncul (mis. 0.50 = 50%)
GUIDE_THRESHOLD = float(os.environ.get("GUIDE_THRESHOLD", "0.50"))

# Optional: suppress TF noisy logs
os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")  # 0=all, 1=info, 2=warning, 3=error

# Lazy load
_model = None
_class_names = None


def _must_exist(path: str, label: str) -> Path:
    p = Path(path)
    if not p.exists():
        raise FileNotFoundError(
            f"{label} tidak ditemukan: '{path}'.\n"
            f"Pastikan file ada di repo Space dan path-nya benar.\n"
            f"Struktur yang disarankan:\n"
            f"  export/best_model.keras\n"
            f"  export/class_names.json"
        )
    return p


def load_assets():
    """Load model and class names once (lazy)."""
    global _model, _class_names

    if _model is None:
        _must_exist(MODEL_PATH, "Model file")
        _model = tf.keras.models.load_model(MODEL_PATH, compile=False)

    if _class_names is None:
        _must_exist(CLASSES_PATH, "Class names file")
        with open(CLASSES_PATH, "r", encoding="utf-8") as f:
            data = json.load(f)

        # class_names.json boleh berupa:
        # - list: ["kelas1","kelas2",...]
        # - dict: {"0":"kelas1","1":"kelas2",...}
        if isinstance(data, dict):
            _class_names = [v for k, v in sorted(data.items(), key=lambda kv: int(kv[0]))]
        elif isinstance(data, list):
            _class_names = data
        else:
            raise ValueError("class_names.json harus berupa list atau dict index->label.")

        if len(_class_names) == 0:
            raise ValueError("class_names.json kosong.")

    return _model, _class_names


def preprocess(img: Image.Image) -> np.ndarray:
    img = img.convert("RGB").resize((IMG_SIZE, IMG_SIZE))
    arr = np.asarray(img, dtype=np.float32) / 255.0
    return np.expand_dims(arr, axis=0)  # (1,H,W,3)


# ============================================================
# Panduan tindakan (1 panduan saja) — berdasarkan label TOP-1
# ============================================================
def _norm_label(s: str) -> str:
    return (s or "").strip().lower().replace("-", " ").replace("_", " ")


GUIDE_INSECT = """\
## Kerusakan akibat serangga / hama

**Tindakan cepat (prioritas 0–24 jam):**
1. **Isolasi** item (kantong plastik bersih/zip) agar hama tidak menyebar ke koleksi lain.
2. **Jangan disikat keras.** Jika ada serbuk/fragmen, lakukan pembersihan **kering dan sangat lembut** (kuas halus) di area kerja terpisah.
3. **Karantina & catat**: lokasi temuan, tingkat kerusakan, indikasi serangga hidup/telur.

**Stabilisasi & penanganan:**
- Jika tersedia fasilitas, pertimbangkan **perlakuan beku (freezing)** untuk mematikan serangga/telur (sesuai SOP lembaga; bungkus berlapis untuk mencegah kondensasi).
- Setelah bebas hama, lakukan **perbaikan fisik** (penguatan sobekan, perataan, perbaikan jilid) oleh petugas kompeten.

**Pencegahan berulang (lingkungan):**
- Housekeeping: vakum rak/area, buang sumber makanan (debu/organik).
- Pantau perangkap serangga, cek titik lembap, perbaiki ventilasi.
- Jaga lingkungan simpan stabil (kelembapan tidak tinggi, sirkulasi baik).
"""

GUIDE_WATER_MOLD = """\
## Kerusakan akibat air / lembap / mold (jamur)

**Tindakan cepat (prioritas 0–24 jam):**
1. **Pisahkan** item basah dari yang kering. Hindari menumpuk (mencegah transfer tinta & jamur).
2. Jika **basah**: lakukan **pengeringan darurat** secepatnya.
   - Bentangkan di permukaan bersih, beri alas penyerap.
   - Gunakan **kipas/sirkulasi udara** (hindari udara panas langsung).
3. Jika ada indikasi **mold** (bau apek kuat, bercak berbulu/berdebu): **isolasi** dan gunakan APD minimum (masker/respirator sesuai SOP, sarung tangan).

**Stabilisasi & penanganan:**
- Untuk item sangat basah/banyak: pertimbangkan **freeze-dry / pembekuan** sebagai stabilisasi darurat (sesuai SOP & fasilitas).
- Pembersihan mold dilakukan **kering** (vakum HEPA/kuas lembut) di area terkontrol; jangan menggosok agresif.
- Setelah kering stabil, evaluasi distorsi/tinta luntur/halaman menempel—rujuk konservator bila parah.

**Kontrol lingkungan (kunci pencegahan jamur):**
- Turunkan risiko dengan mengurangi kelembapan dan memperbaiki sumber kebocoran.
- Pastikan ruang simpan berventilasi baik; monitoring suhu–RH berkala.
- Bersihkan area terdampak dan inspeksi rak sekitar.
"""

DEFAULT_GUIDE = """\
## Panduan umum (kepercayaan model belum cukup tinggi)

**Langkah aman:**
1. **Stabilisasi dulu**: minimalkan sentuhan, jangan ditekan, simpan rata.
2. **Dokumentasi**: foto, lokasi, tanggal, gejala kerusakan.
3. **Isolasi bila ragu** (bau apek, bercak, serbuk): cegah penyebaran ke koleksi lain.
4. **Rujuk SOP institusi/konservator** untuk tindakan lanjut.

Ambang panduan spesifik saat ini: **{thr:.0%}** (GUIDE_THRESHOLD).
"""


def guide_for_top1(label_top1: str, prob_top1: float) -> str:
    """
    Keluarkan 1 panduan saja:
    - Jika prob_top1 >= threshold: panduan sesuai kelas top-1
    - Jika tidak: panduan umum
    """
    if prob_top1 < GUIDE_THRESHOLD:
        return DEFAULT_GUIDE.format(thr=GUIDE_THRESHOLD)

    lab = _norm_label(label_top1)

    # Cocokkan label top-1 ke kategori panduan (sesuaikan keyword sesuai class_names Anda)
    insect_keys = ["serangga", "insect", "hama", "rayap", "termite", "kutu", "booklice"]
    water_mold_keys = ["air", "water", "basah", "banjir", "lembap", "humidity", "mold", "jamur", "kapang"]

    if any(k in lab for k in insect_keys):
        return GUIDE_INSECT

    if any(k in lab for k in water_mold_keys):
        return GUIDE_WATER_MOLD

    # Kalau top-1 tinggi tapi tidak termasuk dua kategori di atas
    return """\
## Rekomendasi umum untuk kelas terdeteksi

- **Prioritaskan stabilisasi** (minimalkan penanganan, simpan rata, dokumentasi).
- Ikuti **SOP konservasi** sesuai tipe kerusakan pada label.
- Bila ada risiko kontaminasi (bau apek/bercak/serbuk), lakukan **isolasi**.

Catatan: Label terdeteksi belum dipetakan ke panduan khusus.
"""


def predict_image(image: Image.Image):
    """
    Return:
      1) dict label->prob (untuk gr.Label top-k)
      2) markdown panduan (hanya 1)
    """
    model, class_names = load_assets()

    x = preprocess(image)
    probs = model.predict(x, verbose=0)[0]  # (num_classes,)

    probs = np.asarray(probs, dtype=np.float32).reshape(-1)
    k = int(min(TOP_K, probs.shape[0], len(class_names)))
    topk_idx = probs.argsort()[-k:][::-1]

    out = {}
    top_labels = []
    top_probs = []
    for idx in topk_idx:
        idx = int(idx)
        label = class_names[idx] if idx < len(class_names) else f"class_{idx}"
        p = float(probs[idx])
        out[label] = p
        top_labels.append(label)
        top_probs.append(p)

    label_top1 = top_labels[0]
    prob_top1 = top_probs[0]

    header = (
        f"### Rekomendasi tindakan berbasis prediksi (Top-{k})\n"
        f"- **Prediksi utama:** **{label_top1}** (**{prob_top1:.0%}**)\n"
        f"- Ambang panduan spesifik: **{GUIDE_THRESHOLD:.0%}**\n\n"
    )

    guide_md = guide_for_top1(label_top1, prob_top1)
    return out, header + guide_md


def build_demo():
    title = "Preservation Classifier"
    desc = "Upload an image of a document/paper fragment. The model predicts the class (top-k) + provides ONE preservation guidance based on the top-1 probability threshold."

    with gr.Blocks() as demo:
        gr.Markdown(f"# {title}\n{desc}")

        with gr.Row():
            inp = gr.Image(type="pil", label="Upload image")
            out_label = gr.Label(num_top_classes=TOP_K, label="Predictions")

        out_guide = gr.Markdown(label="Panduan Tindakan Preservasi")

        with gr.Row():
            btn_clear = gr.Button("Clear")
            btn_submit = gr.Button("Submit", variant="primary")

        btn_submit.click(fn=predict_image, inputs=inp, outputs=[out_label, out_guide])
        btn_clear.click(fn=lambda: (None, None, ""), inputs=None, outputs=[inp, out_label, out_guide])

    return demo


demo = build_demo()

if __name__ == "__main__":
    try:
        load_assets()
    except Exception as e:
        print("Asset load warning:", e)
    demo.launch(server_name="0.0.0.0", server_port=7860)