app.py

import os
import json
import re
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers
from PIL import Image
import io
import asyncio
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
import uvicorn
from huggingface_hub import hf_hub_download, InferenceClient

# Module de traduction multi-moteurs (même dossier)
from translator import translate, translate_fields

# ══════════════════════════════════════════════════════════════════════════════
# CONFIGURATION
# ══════════════════════════════════════════════════════════════════════════════

HF_TOKEN = os.environ.get("HF_TOKEN")
REPO_ID  = "Expendadeur/agro-bio-models"

CONF_REFUSE  = 0.15
CONF_PRUDENT = 0.50
CONF_MOYEN   = 0.75

LANG_DISPLAY = {
    "fr": "Français",
    "en": "English",
    "sw": "Swahili",
    "rn": "Rundi",
}

KNOWN_PLANTS = {
    "banana", "cassava", "cauliflower", "cotton", "guava", "jute",
    "maize", "mango", "papaya", "potato", "rice",
    "sugarcane", "tea", "tomato", "wheat"
}

EXPERT_PROVIDERS = [
    ("groq",           "meta-llama/Llama-3.3-70B-Instruct"),
    ("sambanova",      "meta-llama/Llama-3.3-70B-Instruct"),
    ("novita",         "meta-llama/Llama-3.1-8B-Instruct"),
    ("novita",         "meta-llama/Meta-Llama-3-8B-Instruct"),
    ("featherless-ai", "meta-llama/Meta-Llama-3-8B-Instruct"),
    ("featherless-ai", "HuggingFaceH4/zephyr-7b-beta"),
    ("cerebras",       "meta-llama/Llama-3.1-8B-Instruct"),
]

# ── Labels d'en-tête de section par langue ────────────────────────────────────
# Utilisés pour structurer l'expert_online_advice avec nom maladie + symptômes
# en tête, quelle que soit la langue.
SECTION_LABELS = {
    "fr": {
        "disease":  "MALADIE DETECTEE",
        "symptoms": "SYMPTOMES",
        "advice":   "CONSEILS EXPERT",
        "fallback": "CONSEILS DE BASE",
    },
    "en": {
        "disease":  "DETECTED DISEASE",
        "symptoms": "SYMPTOMS",
        "advice":   "EXPERT ADVICE",
        "fallback": "BASIC ADVICE",
    },
    "sw": {
        "disease":  "UGONJWA ULIOTAMBULIWA",
        "symptoms": "DALILI",
        "advice":   "USHAURI WA MTAALAMU",
        "fallback": "USHAURI WA MSINGI",
    },
    "rn": {
        "disease":  "INDWARA YABONETSE",
        "symptoms": "IBIMENYETSO",
        "advice":   "INAMA Z INZOBERE",
        "fallback": "INAMA ZA MBERE",
    },
}

def get_advice_instruction(conf: float) -> str:
    if conf < CONF_PRUDENT:
        return (
            "Confiance faible (moins de 50%). "
            "Donne UNIQUEMENT des conseils d hygiene et d observation. "
            "Aucun produit chimique. Entre 4 et 6 conseils."
        )
    elif conf < CONF_MOYEN:
        return (
            "Confiance moyenne (50 a 75%). "
            "Methodes biologiques et naturelles uniquement. "
            "Aucun fongicide ou pesticide chimique. "
            "Entre 6 et 9 conseils detailles."
        )
    else:
        return (
            "Confiance elevee (plus de 75%). "
            "Donne TOUS les conseils utiles sans limite : "
            "traitements chimiques si necessaires avec doses precises, "
            "calendrier d application, precautions de securite, "
            "conseils post-traitement, stockage, rotation culturale. "
            "Minimum 10 conseils complets et detailles."
        )

SAFETY_MESSAGES = {
    "fr": {
        "refuse": "IMAGE NON RECONNUE ({:.0f}%). Reprenez la photo : une feuille ou un fruit malade, bien eclaire, cadre de pres.",
        "low":    "PISTE A CONFIRMER ({:.0f}%). Observation recommandee avant tout traitement.",
        "medium": "DIAGNOSTIC PROBABLE ({:.0f}%). Commencez par des methodes naturelles.",
        "high":   "DIAGNOSTIC FIABLE ({:.0f}%). Protocole de traitement recommande.",
    },
    "en": {
        "refuse": "IMAGE NOT RECOGNIZED ({:.0f}%). Retake the photo: a diseased leaf or fruit, well-lit, close-up.",
        "low":    "UNCONFIRMED ({:.0f}%). Observe carefully before any treatment.",
        "medium": "PROBABLE DIAGNOSIS ({:.0f}%). Start with natural methods only.",
        "high":   "RELIABLE DIAGNOSIS ({:.0f}%). Full treatment protocol recommended.",
    },
    "sw": {
        "refuse": "PICHA HAIJATAMBULIWA ({:.0f}%). Piga tena picha ya jani au tunda lenye ugonjwa, mwanga mzuri, karibu.",
        "low":    "HAIJATHIBITISHWA ({:.0f}%). Angalia kwa makini kabla ya matibabu yoyote.",
        "medium": "UTAMBUZI UNAOWEZEKANA ({:.0f}%). Anza na mbinu za asili peke yake.",
        "high":   "UTAMBUZI WA KUAMINIKA ({:.0f}%). Itifaki kamili ya matibabu inashauriwa.",
    },
    "rn": {
        "refuse": "IFOTO NTIYABONETSE ({:.0f}%). Fata ifoto nshya y ikibabi cyangwa imbuto irwaye, mu mucyo mzuri.",
        "low":    "BISABWA KWEMEZWA ({:.0f}%). Kurikirana neza indwara mbere yo gukoresha umuti.",
        "medium": "ISANO IRASHOBOKA ({:.0f}%). Tangira gukoresha imiti ya kamere gusa.",
        "high":   "ISANO YEMEJWE ({:.0f}%). Gukurikira uburyo bwo kuvura burasabwa.",
    },
}

FALLBACK_INTRO = {
    "fr": "Service expert indisponible. Conseils de base :",
    "en": "Expert service unavailable. Basic advice:",
    "sw": "Huduma ya mtaalamu haipatikani. Ushauri wa msingi:",
    "rn": "Serivisi y inzobere ntiboneka. Inama za mbere:",
}

# ══════════════════════════════════════════════════════════════════════════════
# HELPERS
# ══════════════════════════════════════════════════════════════════════════════

def resolve_lang(code: str) -> str:
    code = (code or "fr").lower().strip().split("-")[0]
    return code if code in LANG_DISPLAY else "fr"


def get_safety_msg(lang_code: str, level: str, conf: float) -> str:
    msgs = SAFETY_MESSAGES.get(lang_code, SAFETY_MESSAGES["fr"])
    return msgs.get(level, msgs["low"]).format(conf * 100)


def get_section_label(lang_code: str, key: str) -> str:
    labels = SECTION_LABELS.get(lang_code, SECTION_LABELS["fr"])
    return labels.get(key, SECTION_LABELS["fr"][key])


def clean_text(text: str) -> str:
    if not text:
        return ""
    pattern = re.compile(
        "[" +
        u"\U0001F600-\U0001F64F" +
        u"\U0001F300-\U0001F5FF" +
        u"\U0001F680-\U0001F9FF" +
        u"\U00002702-\U000027B0" +
        u"\U000024C2-\U0001F251" +
        u"\U0001F900-\U0001F9FF" +
        u"\U00002600-\U000026FF" +
        u"\U00002700-\U000027BF" +
        "]+", flags=re.UNICODE
    )
    return pattern.sub("", text).strip()


def to_bullets(text: str) -> str:
    """Convertit n importe quel texte en tirets verticaux propres."""
    if not text:
        return ""
    text = clean_text(text)

    if "\n" in text and "-" in text:
        lines = []
        for line in text.split("\n"):
            line = line.strip()
            if not line:
                continue
            if not line.startswith("-"):
                line = "- " + line
            lines.append(line)
        return "\n".join(lines)

    if " - " in text:
        parts = re.split(r'\s+-\s+', text)
        lines = []
        for p in parts:
            p = p.strip().lstrip("-").strip()
            if p and len(p) > 4:
                lines.append("- " + p)
        if len(lines) >= 2:
            return "\n".join(lines)

    sentences = re.split(r'(?<=[.!?])\s+', text)
    if len(sentences) >= 2:
        lines = []
        for s in sentences:
            s = s.strip()
            if s and len(s) > 8:
                lines.append("- " + s)
        if lines:
            return "\n".join(lines)

    return "- " + text.strip()


def ood_message(lang_code: str, plant_name: str) -> str:
    plants = ", ".join(sorted(KNOWN_PLANTS))
    msgs = {
        "fr": f"Plante '{plant_name}' non reconnue. Cultures supportees : {plants}.",
        "en": f"Plant '{plant_name}' not recognized. Supported crops: {plants}.",
        "sw": f"Mmea '{plant_name}' haujatambuliwa. Mazao yanayoungwa mkono: {plants}.",
        "rn": f"Ikimera '{plant_name}' ntikizwi. Ibimera bizwi: {plants}.",
    }
    return msgs.get(lang_code, msgs["fr"])


def build_ref_words(
    treatment_fr: str,
    prevention_fr: str,
    treatment_translated: str,
    prevention_translated: str,
    cause_fr: str = "",
    cause_translated: str = "",
) -> set:
    """
    Construit un ensemble de mots de référence dans TOUTES les langues disponibles
    (français + langue cible) pour un filtrage anti-doublon efficace quelle que
    soit la langue de la réponse de l IA.
    """
    combined = " ".join([
        treatment_fr,
        prevention_fr,
        cause_fr,
        treatment_translated,
        prevention_translated,
        cause_translated,
    ]).lower()
    return set(combined.split())


def filter_duplicate_lines(
    ia_lines: list,
    ref_words: set,
    overlap_threshold: float = 0.70,
) -> list:
    """
    Filtre les lignes de l IA dont le contenu chevauche trop la KB de référence.
    Fonctionne dans n importe quelle langue car ref_words contient les deux langues.
    """
    filtered = []
    for line in ia_lines:
        line = line.strip()
        if not line:
            continue
        line_words = set(line.lower().split())
        if not line_words:
            continue
        overlap = len(line_words & ref_words) / len(line_words)
        if overlap < overlap_threshold:
            filtered.append(line)
        else:
            print(f"Doublon filtre ({overlap:.0%} overlap) : {line[:60]}")
    return filtered


def build_expert_online_advice(
    safety_msg: str,
    disease_name: str,
    symptoms_block: str,
    advice_block: str,
    lang_code: str,
    is_fallback: bool = False,
) -> str:
    """
    Construit le champ expert_online_advice avec une structure fixe et lisible :
        [safety_msg]
        === MALADIE DETECTEE ===
        [disease_name]
        === SYMPTOMES ===
        [symptoms_block]
        === CONSEILS EXPERT / CONSEILS DE BASE ===
        [advice_block]
    Cette structure est TOUJOURS respectée, que l IA soit online ou offline,
    afin que l utilisateur sache immédiatement à quelle maladie appartiennent
    les informations affichées.
    """
    lbl_disease  = get_section_label(lang_code, "disease")
    lbl_symptoms = get_section_label(lang_code, "symptoms")
    lbl_advice   = get_section_label(lang_code, "fallback" if is_fallback else "advice")

    parts = [
        safety_msg,
        "",
        f"=== {lbl_disease} ===",
        disease_name,
        "",
        f"=== {lbl_symptoms} ===",
        symptoms_block,
        "",
        f"=== {lbl_advice} ===",
        advice_block,
    ]
    return "\n".join(parts)


# ══════════════════════════════════════════════════════════════════════════════
# EXPERT IA — FLUX COMPLET
# ══════════════════════════════════════════════════════════════════════════════
#
# FLUX :
#   1. L IA genere librement depuis la maladie détectée :
#      - D abord : NOM DE LA MALADIE + SYMPTOMES (section obligatoire)
#      - Ensuite : conseils/traitement/prevention complémentaires
#   2. On extrait la section symptômes et la section conseils séparément.
#   3. On filtre les doublons des conseils vs KB (FR + langue cible).
#   4. build_expert_online_advice assemble la réponse finale structurée.
#
# ══════════════════════════════════════════════════════════════════════════════

# Marqueur de séparation entre symptômes et conseils dans la réponse IA
_SYMPTOMS_MARKER = "###SYMPTOMES###"
_ADVICE_MARKER   = "###CONSEILS###"


def call_expert_ia(
    disease_name: str,
    plant_name: str,
    conf: float,
    tone: str,
    lang_code: str,
    treatment_fr: str,
    prevention_fr: str,
    cause_fr: str,
    treatment_translated: str,
    prevention_translated: str,
    cause_translated: str,
) -> dict | None:
    """
    Retourne un dict {"symptoms": str, "advice": str} ou None si échec.
    L IA génère deux blocs distincts :
      - symptoms : symptômes visibles de la maladie (1 tiret par ligne)
      - advice   : conseils/traitement/prevention complémentaires (1 tiret par ligne)
    Les deux blocs sont séparés par _ADVICE_MARKER dans la réponse brute.
    """

    advice_instruction = get_advice_instruction(conf)
    lang_display       = LANG_DISPLAY.get(lang_code, "Français")

    lang_note = ""
    if lang_code == "rn":
        lang_note = (
            " IMPORTANT : Tu reponds en Kirundi (appele aussi Rundi),"
            " la langue nationale du Burundi."
            " Ce n est pas du Swahili."
        )

    system_prompt = (
        f"Tu es l Expert Agronome AGRO DIAG au Burundi."
        f" Reponds UNIQUEMENT en {lang_display}."
        f" Sans emojis. Sans icones. Sans introduction. Sans conclusion.{lang_note}"
        f" Niveau de certitude : {tone}."
        f" Conseils adaptes aux petits agriculteurs burundais."
        f" Format strict decrit ci-dessous."
        f" Reponds UNIQUEMENT avec les deux blocs demandes, rien d autre."
    )

    # Le prompt impose deux blocs clairement séparés par le marqueur.
    # Bloc 1 — Symptômes : ce que l agriculteur voit sur la plante.
    # Bloc 2 — Conseils : enrichissement calibré selon le niveau de confiance.
    user_prompt = (
        f"Maladie detectee : {disease_name} sur {plant_name}.\n"
        f"Niveau de confiance du diagnostic : {conf*100:.0f}%.\n\n"

        f"Genere exactement DEUX blocs en {lang_display} :\n\n"

        f"BLOC 1 — Symptomes visibles :\n"
        f"Ecris le marqueur exactement : {_ADVICE_MARKER}\n"
        f"Puis liste en tirets (-) les symptomes caracteristiques de {disease_name} "
        f"tels qu ils apparaissent sur la plante (feuilles, tiges, fruits, racines).\n"
        f"Entre 3 et 5 symptomes. Un tiret par ligne.\n\n"

        f"BLOC 2 — Conseils et traitements :\n"
        f"Ecris le marqueur exactement : {_ADVICE_MARKER}\n"
        f"Puis liste en tirets (-) les conseils pratiques adaptes au niveau de confiance.\n"
        f"{advice_instruction}\n\n"

        f"FORMAT FINAL ATTENDU (respecte exactement) :\n"
        f"{_SYMPTOMS_MARKER}\n"
        f"- symptome 1\n"
        f"- symptome 2\n"
        f"...\n"
        f"{_ADVICE_MARKER}\n"
        f"- conseil 1\n"
        f"- conseil 2\n"
        f"...\n"
    )

    # Mots de référence KB dans les deux langues pour filtrage anti-doublon
    ref_words = build_ref_words(
        treatment_fr, prevention_fr,
        treatment_translated, prevention_translated,
        cause_fr, cause_translated,
    )

    for provider, model_id in EXPERT_PROVIDERS:
        try:
            c = InferenceClient(api_key=HF_TOKEN, provider=provider)
            response = c.chat.completions.create(
                model=model_id,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user",   "content": user_prompt},
                ],
                max_tokens=1800,
                temperature=0.2,
            )
            raw = clean_text(response.choices[0].message.content.strip())

            # ── Extraction des deux blocs ──────────────────────────────────────
            symptoms_block = ""
            advice_block   = ""

            if _SYMPTOMS_MARKER in raw and _ADVICE_MARKER in raw:
                # Cas idéal : les deux marqueurs sont présents
                idx_sym = raw.index(_SYMPTOMS_MARKER)
                idx_adv = raw.index(_ADVICE_MARKER)
                symptoms_raw = raw[idx_sym + len(_SYMPTOMS_MARKER):idx_adv].strip()
                advice_raw   = raw[idx_adv + len(_ADVICE_MARKER):].strip()
            elif _ADVICE_MARKER in raw:
                # Seulement le marqueur conseils — on met tout avant en symptômes
                idx_adv = raw.index(_ADVICE_MARKER)
                symptoms_raw = raw[:idx_adv].strip()
                advice_raw   = raw[idx_adv + len(_ADVICE_MARKER):].strip()
            else:
                # Aucun marqueur — on coupe à mi-chemin (fallback de parsing)
                lines_all    = [l.strip() for l in raw.split("\n") if l.strip()]
                mid          = max(1, len(lines_all) // 3)
                symptoms_raw = "\n".join(lines_all[:mid])
                advice_raw   = "\n".join(lines_all[mid:])

            symptoms_lines = [l.strip() for l in symptoms_raw.split("\n") if l.strip()]
            advice_lines   = [l.strip() for l in advice_raw.split("\n")   if l.strip()]

            # Normalisation tirets
            symptoms_lines = [l if l.startswith("-") else "- " + l for l in symptoms_lines]
            advice_lines   = [l if l.startswith("-") else "- " + l for l in advice_lines]

            # ── Filtrage anti-doublon des CONSEILS uniquement ─────────────────
            # (les symptômes ne sont pas filtrés — ils sont toujours pertinents)
            filtered_advice = filter_duplicate_lines(advice_lines, ref_words, overlap_threshold=0.70)

            nb_sym = len(symptoms_lines)
            nb_adv = len(filtered_advice)

            if nb_sym < 1:
                print(f"Aucun symptome extrait — {provider}, on tente le suivant")
                continue
            if nb_adv < 2:
                print(f"Trop peu de conseils uniques ({nb_adv}) — {provider}, on tente le suivant")
                continue

            symptoms_block = "\n".join(symptoms_lines)
            advice_block   = "\n".join(filtered_advice)

            print(
                f"Expert IA OK ({lang_display}, {nb_sym} symptomes, {nb_adv} conseils) "
                f"— {provider}/{model_id.split('/')[-1]}"
            )
            return {"symptoms": symptoms_block, "advice": advice_block}

        except Exception as e:
            print(f"Expert IA ECHEC — {provider} : {str(e)[:80]}")
            continue

    return None


# ══════════════════════════════════════════════════════════════════════════════
# CUSTOM LAYER
# ══════════════════════════════════════════════════════════════════════════════

@tf.keras.utils.register_keras_serializable(package="Custom")
class GeMPooling(layers.Layer):
    def __init__(self, p=3.0, **kwargs):
        super().__init__(**kwargs)
        self.p = tf.Variable(p, trainable=True, dtype=tf.float32, name="gem_p")

    def call(self, x):
        x = tf.cast(x, tf.float32)
        x = tf.clip_by_value(x, 1e-6, tf.reduce_max(x))
        x = tf.pow(x, self.p)
        x = tf.reduce_mean(x, axis=[1, 2])
        return tf.pow(x, 1.0 / self.p)

    def get_config(self):
        cfg = super().get_config()
        cfg.update({"p": float(self.p.numpy())})
        return cfg


# ══════════════════════════════════════════════════════════════════════════════
# FASTAPI
# ══════════════════════════════════════════════════════════════════════════════

app = FastAPI(title="Agro Diag Burundi — API V14")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# ── Chargement labels ─────────────────────────────────────────────────────────
try:
    labels_path = hf_hub_download(
        repo_id=REPO_ID, filename="labels.json",
        token=HF_TOKEN, repo_type="model"
    )
    with open(labels_path) as f:
        labels = json.load(f)
    print(f"Labels : {len(labels)} classes")
except Exception as e:
    print(f"Labels ECHEC : {e}")
    labels = []

# ── Chargement poids optimaux ─────────────────────────────────────────────────
try:
    weights_path = hf_hub_download(
        repo_id=REPO_ID, filename="online_weights.json",
        token=HF_TOKEN, repo_type="model"
    )
    with open(weights_path) as f:
        weights_data = json.load(f)
    model_names    = weights_data["model_names"]
    optuna_weights = weights_data["weights_optuna"]
    weights_map    = dict(zip(model_names, optuna_weights))
    print(f"Poids Optuna charges : {weights_map}")
except Exception as e:
    print(f"Poids ECHEC : {e} — poids egaux utilises")
    weights_map = {
        "EfficientNetV2M": 0.25,
        "ResNet50":        0.25,
        "MobileNetV3L":    0.25,
        "ConvNeXtBase":    0.25,
    }

# ── Configuration modeles ─────────────────────────────────────────────────────
MODEL_CONFIG = {
    "EfficientNetV2M": {
        "filename":   "online_efficientnetv2m.keras",
        "size":       (260, 260),
        "preprocess": tf.keras.applications.efficientnet_v2.preprocess_input,
    },
    "ResNet50": {
        "filename":   "online_resnet50.keras",
        "size":       (260, 260),
        "preprocess": tf.keras.applications.resnet.preprocess_input,
    },
    "MobileNetV3L": {
        "filename":   "online_mobilenetv3l.keras",
        "size":       (224, 224),
        "preprocess": tf.keras.applications.mobilenet_v2.preprocess_input,
    },
    "ConvNeXtBase": {
        "filename":   "online_convnextbase.keras",
        "size":       (384, 384),
        "preprocess": tf.keras.applications.convnext.preprocess_input,
    },
}

for name in MODEL_CONFIG:
    MODEL_CONFIG[name]["weight"] = weights_map.get(name, 0.25)

loaded_models = {}
for name, cfg in MODEL_CONFIG.items():
    try:
        path = hf_hub_download(
            repo_id=REPO_ID, filename=cfg["filename"],
            token=HF_TOKEN, repo_type="model"
        )
        loaded_models[name] = tf.keras.models.load_model(
            path, custom_objects={"GeMPooling": GeMPooling}
        )
        print(f"{name} pret (poids={cfg['weight']:.4f})")
    except Exception as e:
        print(f"{name} ECHEC : {e}")

# ── Chargement base de connaissances ──────────────────────────────────────────
try:
    kb_path = hf_hub_download(
        repo_id=REPO_ID, filename="disease_kb.json",
        token=HF_TOKEN, repo_type="model"
    )
    with open(kb_path, encoding="utf-8") as f:
        disease_kb = json.load(f)
    print(f"KB : {len(disease_kb)} entrees")
except Exception as e:
    print(f"KB ECHEC : {e}")
    disease_kb = {}


def preprocess_image(image, target_size, preprocess_fn):
    if image.mode != "RGB":
        image = image.convert("RGB")
    image = image.resize(target_size)
    arr = np.expand_dims(np.array(image).astype(np.float32), axis=0)
    return preprocess_fn(arr)


# ══════════════════════════════════════════════════════════════════════════════
# ENDPOINT /predict
# ══════════════════════════════════════════════════════════════════════════════

@app.post("/predict")
async def predict(request: Request):
    try:
        contents = await request.body()

        if not contents:
            return JSONResponse(status_code=400, content={"error": "Corps vide."})
        if len(contents) < 3 * 1024:
            return JSONResponse(status_code=400,
                content={"error": "Image trop petite (minimum 3 KB)."})

        lang_code = resolve_lang(
            request.query_params.get("lang")
            or request.headers.get("x-language")
            or "fr"
        )

        try:
            image = Image.open(io.BytesIO(contents))
            image.verify()
            image = Image.open(io.BytesIO(contents))
        except Exception:
            return JSONResponse(status_code=400,
                content={"error": "Image invalide ou corrompue."})

        w, h = image.size
        if w < 50 or h < 50:
            return JSONResponse(status_code=400,
                content={"error": f"Image trop petite ({w}x{h}px). Minimum 50x50."})

        # ── Ensemble — inference ───────────────────────────────────────────────
        ensemble_probs = np.zeros(len(labels))
        total_weight   = 0.0
        for name, model in loaded_models.items():
            cfg        = MODEL_CONFIG[name]
            input_data = preprocess_image(image, cfg["size"], cfg["preprocess"])
            probs      = model.predict(input_data, verbose=0)[0]
            ensemble_probs += probs * cfg["weight"]
            total_weight   += cfg["weight"]
        if total_weight > 0:
            ensemble_probs /= total_weight

        top_idx    = int(np.argmax(ensemble_probs))
        top_class  = labels[top_idx]
        top_conf   = float(ensemble_probs[top_idx])
        plant_name = top_class.split("_")[0]

        # ── KB locale (source primaire garantie) ──────────────────────────────
        kb              = disease_kb.get(top_class, {})
        disease_name_fr = kb.get("display_name", top_class.replace("_", " ").title())
        cause_fr        = kb.get("cause",      "Information non disponible.")
        treatment_fr    = kb.get("treatment",  "Consultez un expert agronome.")
        prevention_fr   = kb.get("prevention", "Maintenir une bonne hygiene de culture.")
        # Symptômes KB (fallback si l IA est offline)
        symptoms_kb_fr  = kb.get("symptoms",   cause_fr)

        # ── Garde-fou : confiance trop faible ─────────────────────────────────
        if top_conf < CONF_REFUSE:
            msg = get_safety_msg(lang_code, "refuse", top_conf)
            return JSONResponse({
                "label":                top_class,
                "confidence":           top_conf,
                "disease_name":         "—",
                "lang":                 lang_code,
                "cause":                msg,
                "treatment":            "—",
                "prevention":           "—",
                "expert_online_advice": msg,
                "warning":              "confidence_too_low",
            })

        # ── Garde-fou : plante hors domaine ───────────────────────────────────
        if plant_name not in KNOWN_PLANTS:
            msg = ood_message(lang_code, plant_name)
            return JSONResponse({
                "label":                top_class,
                "confidence":           top_conf,
                "disease_name":         clean_text(disease_name_fr),
                "lang":                 lang_code,
                "cause":                "—",
                "treatment":            "—",
                "prevention":           "—",
                "expert_online_advice": msg,
                "warning":              "out_of_domain",
            })

        # ── Niveau de confiance ────────────────────────────────────────────────
        if top_conf < CONF_PRUDENT:
            tone       = "PRUDENT : hygiene et observation uniquement. Aucun produit chimique."
            safety_msg = get_safety_msg(lang_code, "low", top_conf)
        elif top_conf < CONF_MOYEN:
            tone       = "CONSEILLER EXPERT : methodes biologiques et naturelles uniquement."
            safety_msg = get_safety_msg(lang_code, "medium", top_conf)
        else:
            tone       = "EXPERT CONFIRME : protocole complet, produits chimiques autorises si necessaires."
            safety_msg = get_safety_msg(lang_code, "high", top_conf)

        # ── Traduction KB ──────────────────────────────────────────────────────
        loop = asyncio.get_event_loop()

        fields_fr = {
            "disease_name": disease_name_fr,
            "cause":        cause_fr,
            "treatment":    treatment_fr,
            "prevention":   prevention_fr,
            "symptoms_kb":  symptoms_kb_fr,
        }

        translated = await loop.run_in_executor(
            None,
            lambda: translate_fields(fields_fr, lang_code)
        )

        disease_name_t  = clean_text(translated.get("disease_name", disease_name_fr))
        cause_t         = to_bullets(clean_text(translated.get("cause",       cause_fr)))
        treatment_t     = to_bullets(clean_text(translated.get("treatment",   treatment_fr)))
        prevention_t    = to_bullets(clean_text(translated.get("prevention",  prevention_fr)))
        symptoms_kb_t   = to_bullets(clean_text(translated.get("symptoms_kb", symptoms_kb_fr)))

        # ── Expert IA ──────────────────────────────────────────────────────────
        #
        # ETAPE 1 : L IA genere librement depuis la maladie
        #           → bloc symptomes  (nom maladie en tête imposé par le prompt)
        #           → bloc conseils   (enrichissement calibré selon confiance)
        # ETAPE 2 : Filtrage doublon des CONSEILS vs KB (FR + langue cible)
        # ETAPE 3 : build_expert_online_advice assemble la réponse finale :
        #           safety_msg / nom maladie / symptomes / conseils
        #
        ia_result = None
        try:
            ia_result = await asyncio.wait_for(
                loop.run_in_executor(
                    None,
                    lambda: call_expert_ia(
                        disease_name          = disease_name_t,
                        plant_name            = plant_name,
                        conf                  = top_conf,
                        tone                  = tone,
                        lang_code             = lang_code,
                        treatment_fr          = treatment_fr,
                        prevention_fr         = prevention_fr,
                        cause_fr              = cause_fr,
                        treatment_translated  = translated.get("treatment",  ""),
                        prevention_translated = translated.get("prevention", ""),
                        cause_translated      = translated.get("cause",      ""),
                    )
                ),
                timeout=60.0
            )
        except asyncio.TimeoutError:
            print("Expert IA timeout 60s — fallback KB locale")

        # ── Construction expert_online_advice ──────────────────────────────────
        #
        # Structure garantie dans TOUS les cas (online ET offline) :
        #
        #   [safety_msg]
        #
        #   === MALADIE DETECTEE ===       ← toujours affiché
        #   [disease_name]
        #
        #   === SYMPTOMES ===              ← toujours affiché
        #   [symptomes IA  OU  symptomes KB]
        #
        #   === CONSEILS EXPERT / CONSEILS DE BASE ===
        #   [conseils IA filtrés  OU  KB traitement]
        #
        if ia_result:
            # Mode online : symptômes et conseils viennent de l IA
            expert_online_advice = build_expert_online_advice(
                safety_msg   = safety_msg,
                disease_name = disease_name_t,
                symptoms_block = ia_result["symptoms"],
                advice_block   = ia_result["advice"],
                lang_code      = lang_code,
                is_fallback    = False,
            )
        else:
            # Mode offline : symptômes KB + traitement KB
            expert_online_advice = build_expert_online_advice(
                safety_msg     = safety_msg,
                disease_name   = disease_name_t,
                symptoms_block = symptoms_kb_t,
                advice_block   = treatment_t,
                lang_code      = lang_code,
                is_fallback    = True,
            )

        return JSONResponse({
            # ── Identification ────────────────────────────────────────────────
            "label":      top_class,
            "confidence": top_conf,
            "lang":       lang_code,

            # ── Nom de la maladie ─────────────────────────────────────────────
            "disease_name": disease_name_t,

            # ── KB locale (source primaire garantie, toujours presente) ───────
            "cause":      cause_t,
            "treatment":  treatment_t,
            "prevention": prevention_t,

            # ── Réponse Expert IA structurée ───────────────────────────────────
            # Commence TOUJOURS par : nom maladie → symptômes → conseils
            # online  : symptômes + conseils générés par l IA, conseils filtrés anti-doublon
            # offline : symptômes KB + traitement KB (fallback garanti)
            "expert_online_advice": expert_online_advice,
        })

    except Exception as e:
        print(f"Erreur predict : {e}")
        return JSONResponse(status_code=500, content={"error": str(e)})


# ══════════════════════════════════════════════════════════════════════════════
# CHAT — assistant vocal agricole (endpoint appelé par le mode online Flutter)
# ══════════════════════════════════════════════════════════════════════════════

CHAT_FALLBACK = {
    "fr": "Service expert indisponible. Consultez un agronome local ou prenez une photo de votre plante pour un diagnostic precis.",
    "en": "Expert service unavailable. Please consult a local agronomist or take a photo of your plant for an accurate diagnosis.",
    "sw": "Huduma ya mtaalamu haipatikani. Wasiliana na mtaalamu wa kilimo wa karibu au piga picha ya mmea wako kwa utambuzi sahihi.",
    "rn": "Serivisi y inzobere ntiboneka. Baza inzobere y uburimyi yo hafi canke fata ifoto y ikimera cawe kugira isano itomoye.",
}


def search_kb_context(query: str) -> tuple[str, list[str], dict]:
    """Recherche dans la KB locale les entrees pertinentes par score de pertinence."""
    q = query.lower()
    
    # Liste basique de mots vides à ignorer
    stopwords = {"sont", "avec", "dans", "pour", "vous", "avez", "nous", "elles", "ils", "cette", "ceux", "celles", "votre", "notre", "comme", "quand", "plus", "très", "tout", "tous"}
    
    words_raw = re.findall(r'\b\w+\b', q)
    query_words = [w for w in words_raw if len(w) > 3 and w not in stopwords]
    
    if not query_words:
        return "", [], {}

    scores = []
    
    for key, entry in disease_kb.items():
        name      = entry.get("display_name", "").lower()
        cause     = entry.get("cause",      "").lower()
        treatment = entry.get("treatment",  "").lower()
        prev      = entry.get("prevention", "").lower()
        combined  = f"{name} {cause} {treatment} {prev}"
        
        score = sum(1 for w in query_words if w in combined)
        # Bonus si le mot exact du nom de la maladie/plante est présent
        if any(w in name for w in query_words):
            score += 2
            
        if score > 0:
            scores.append((score, key, entry))
            
    # Trier par score décroissant
    scores.sort(key=lambda x: x[0], reverse=True)
    top_matches = scores[:3]
    
    if not top_matches:
        return "", [], {}
        
    results_text = []
    found_diseases = []
    first_match = top_matches[0][2]
    
    for score, key, entry in top_matches:
        d_name = entry.get('display_name', key)
        found_diseases.append(d_name)
        results_text.append(
            f"- {d_name} : "
            f"traitement={entry.get('treatment', '')[:120]} | "
            f"prevention={entry.get('prevention', '')[:80]}"
        )

    return "\n".join(results_text), found_diseases, first_match


def call_chat_ia(query: str, lang_code: str, context: str, diseases: list[str]) -> str | None:
    """Appelle le LLM pour repondre a une question agricole en langage naturel."""
    lang_display = LANG_DISPLAY.get(lang_code, "Français")

    lang_note = ""
    if lang_code == "rn":
        lang_note = (
            " IMPORTANT : Tu reponds en Kirundi (appele aussi Rundi),"
            " la langue nationale du Burundi. Ce n est pas du Swahili."
        )

    context_block = ""
    if context:
        context_block = (
            f"\nBase de donnees de reference (utilise si pertinent) :\n{context}\n"
        )

    # Si des maladies ont été trouvées dans la base, on demande au LLM d'indiquer clairement la maladie suspectée
    disease_prompt = ""
    if diseases:
        d_list = ", ".join(diseases)
        disease_prompt = f" Commence ta reponse en citant la ou les maladies suspectees parmi : {d_list}."

    system_prompt = (
        f"Tu es AGRO DIAG, un assistant agricole expert au Burundi."
        f" Reponds UNIQUEMENT en {lang_display}. Sans emojis. Sans introduction formelle."
        f" Sois concis, pratique et direct.{lang_note}"
        f" Conseils adaptes aux petits agriculteurs d Afrique de l Est."
        f" Si la question ne concerne pas l agriculture, reponds poliment"
        f" que tu ne peux repondre qu aux questions agricoles."
        f"{disease_prompt}"
        f"{context_block}"
    )

    user_prompt = (
        f"Question de l agriculteur : {query}\n\n"
        f"Reponds en {lang_display} de facon claire et pratique."
        f" Maximum 5 phrases ou 5 points. Pas de tirets si ce n est pas necessaire."
    )

    for provider, model_id in EXPERT_PROVIDERS:
        try:
            c = InferenceClient(api_key=HF_TOKEN, provider=provider)
            response = c.chat.completions.create(
                model=model_id,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user",   "content": user_prompt},
                ],
                max_tokens=500,
                temperature=0.3,
            )
            raw = clean_text(response.choices[0].message.content.strip())
            if raw and len(raw) > 10:
                print(f"Chat IA OK ({lang_display}) — {provider}/{model_id.split('/')[-1]}")
                return raw
        except Exception as e:
            print(f"Chat IA ECHEC — {provider} : {str(e)[:80]}")
            continue
    return None


@app.post("/chat")
async def chat(request: Request):
    """
    Endpoint pour l assistant vocal Flutter.
    Body JSON : {"query": "...", "lang": "fr|en|sw|rn"}
    Reponse   : {"answer": "..."}
    Timeout interne : 13s (Flutter attend 20s max).
    """
    lang_code = "fr"
    try:
        body      = await request.json()
        query     = (body.get("query") or "").strip()
        lang_code = resolve_lang(body.get("lang") or "fr")

        if not query:
            return JSONResponse({"answer": CHAT_FALLBACK.get(lang_code, CHAT_FALLBACK["fr"]), "diseases": [], "cause": "", "treatment": "", "prevention": ""})

        # Recherche de contexte dans la KB pour enrichir le LLM
        context, found_diseases, first_match = search_kb_context(query)

        loop   = asyncio.get_event_loop()
        
        # ── Traduction des infos KB si on a trouvé une maladie ────────────────
        cause_t = ""
        treatment_t = ""
        prevention_t = ""

        if first_match:
            fields_fr = {
                "cause":      first_match.get("cause", ""),
                "treatment":  first_match.get("treatment", ""),
                "prevention": first_match.get("prevention", "")
            }
            translated = await loop.run_in_executor(
                None,
                lambda: translate_fields(fields_fr, lang_code)
            )
            cause_t      = to_bullets(clean_text(translated.get("cause",      fields_fr["cause"])))
            treatment_t  = to_bullets(clean_text(translated.get("treatment",  fields_fr["treatment"])))
            prevention_t = to_bullets(clean_text(translated.get("prevention", fields_fr["prevention"])))

        # ── Appel Expert IA ───────────────────────────────────────────────────
        answer = None
        try:
            answer = await asyncio.wait_for(
                loop.run_in_executor(
                    None,
                    lambda: call_chat_ia(query, lang_code, context, found_diseases)
                ),
                timeout=13.0   # 13s < timeout Flutter de 20s
            )
        except asyncio.TimeoutError:
            print("Chat IA timeout 13s — fallback")

        if not answer:
            answer = CHAT_FALLBACK.get(lang_code, CHAT_FALLBACK["fr"])

        return JSONResponse({
            "answer": answer,
            "diseases": found_diseases,
            "cause": cause_t,
            "treatment": treatment_t,
            "prevention": prevention_t
        })

    except Exception as e:
        print(f"Erreur chat : {e}")
        return JSONResponse({"answer": CHAT_FALLBACK.get(lang_code, CHAT_FALLBACK["fr"]), "diseases": [], "cause": "", "treatment": "", "prevention": ""})


# ══════════════════════════════════════════════════════════════════════════════
# ENDPOINT GET /
# ══════════════════════════════════════════════════════════════════════════════

@app.get("/")
def read_root():
    return {
        "status":        "online",
        "ensemble":      "V14_Optuna_agro-bio-models",
        "models_source": REPO_ID,
        "models_loaded": list(loaded_models.keys()),
        "classes_count": len(labels),
        "endpoints": {
            "POST /predict": "Analyse d image de plante (multipart ou raw bytes)",
            "POST /chat":    "Assistant vocal agricole (JSON: {query, lang})",
            "GET /":         "Statut de l API",
        },
        "known_plants":  sorted(KNOWN_PLANTS),
        "supported_langs": {
            "fr": "Français — direct (pas de traduction)",
            "en": "English",
            "sw": "Swahili",
            "rn": "Rundi (Kirundi du Burundi)",
        },
        "translation_engines": [
            "1. googletrans==4.0.0-rc1 (Google Translate, gratuit)",
            "2. deep-translator GoogleTranslator (Google, gratuit)",
            "3. deep-translator MyMemoryTranslator (5000 mots/jour gratuit)",
            "4. deep-translator LibreTranslator (open source)",
            "5. Fallback : texte original français conserve",
        ],
        "expert_ia_flow": {
            "etape_1": "L IA genere librement : symptomes + conseils depuis la maladie detectee (sans voir la KB)",
            "etape_2": "Filtrage doublons des CONSEILS vs KB dans les deux langues (FR + langue cible) — seuil 70%",
            "etape_3": "build_expert_online_advice structure la reponse : safety_msg / nom maladie / symptomes / conseils",
            "fallback": "Si IA offline : structure identique avec symptomes KB + traitement KB",
        },
        "expert_online_advice_structure": {
            "1": "[safety_msg]",
            "2": "=== MALADIE DETECTEE === + disease_name",
            "3": "=== SYMPTOMES === + symptomes IA (ou KB si offline)",
            "4": "=== CONSEILS EXPERT === + conseils IA filtres (ou traitement KB si offline)",
        },
        "advice_by_confidence": {
            "moins de 50%": "4 a 6 conseils hygiene uniquement",
            "50 a 75%":     "6 a 9 conseils biologiques",
            "plus de 75%":  "10+ conseils complets sans limite",
        },
        "thresholds": {
            "refuse":  CONF_REFUSE,
            "prudent": CONF_PRUDENT,
            "moyen":   CONF_MOYEN,
        },
        "model_weights": {
            name: MODEL_CONFIG[name]["weight"]
            for name in MODEL_CONFIG
        },
        "expert_providers": [f"{p}/{m.split('/')[-1]}" for p, m in EXPERT_PROVIDERS],
    }


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)