# -*- coding: utf-8 -*-
"""
Assistant RAG (Claude) — panneau flottant pour VizionAyiti2030
===============================================================
Panneau flottant enrichi : menu repliable (filtres thème/secteur/source),
sélecteur de modèle, questions suggérées, « Nouveau chat », langue
(Français / English / Kreyòl), génération de rapports (.docx) et voix kreyòl.
Hyperparamètres -> config.py. Secrets -> .streamlit/secrets.toml.
"""

import os
import glob
import re
import hashlib
import streamlit as st
import rapport
import voice
import translate
import config

ICI = os.path.dirname(os.path.abspath(__file__))
DOCS_DIR = os.path.join(ICI, "documents")
CATALOGUE = os.path.join(ICI, "catalogue_indicateurs_FR.csv")

MODELS = list(dict.fromkeys([config.LLM_MODEL, "claude-sonnet-4-6",
                             "claude-haiku-4-5-20251001", "claude-opus-4-8"]))
LANGS = ["Français", "English", "Kreyòl"]
SUGGESTIONS = [
    "Quel est l'état d'avancement du Plan de relèvement ?",
    "Résume le budget rectificatif 2025–2026.",
    "Quels indicateurs de santé sont disponibles ?",
    "Génère un rapport sur la mortalité maternelle.",
]

MODULE_INFO = {
    "m1": "Suivi du Plan de relèvement (KPIs budget, priorités de transition)",
    "m2": "Analyse hebdomadaire des données en direct (HDX, ACAPS, WFP, brief)",
    "m3": "Veille stratégique & médias (tendances, signaux faibles, alertes)",
    "m4": "Bibliothèque des produits analytiques (SitReps, fiches, rapports)",
    "m5": "Analyse pour les réunions de l'UNCT",
    "m6": "Analyse pour les réunions gouvernementales",
    "m7": "Analyse pour les partenaires de développement",
    "m8": "Intégration de données primaires (réseaux société civile)",
}


# ----------------------------------------------------------------- indexation
def _mtime(p):
    try:
        return os.path.getmtime(p)
    except OSError:
        return 0.0


def _index_paths():
    paths = [CATALOGUE]
    if os.path.isdir(DOCS_DIR):
        paths += sorted(glob.glob(os.path.join(DOCS_DIR, "**", "*"), recursive=True))
    return paths


def _chunks(text, size=config.RAG_CHUNK_SIZE):
    text = re.sub(r"\s+", " ", text).strip()
    return [text[i:i + size] for i in range(0, len(text), size)] if text else []


def _read_document(path):
    ext = path.lower().rsplit(".", 1)[-1] if "." in path else ""
    try:
        if ext in ("txt", "md"):
            with open(path, encoding="utf-8", errors="ignore") as f:
                return f.read()
        if ext == "pdf":
            from pypdf import PdfReader
            return "\n".join((pg.extract_text() or "") for pg in PdfReader(path).pages)
    except Exception:
        return ""
    return ""


@st.cache_resource(show_spinner="Indexation du magasin de données…")
def build_index(sig):
    docs = []
    for key, label in MODULE_INFO.items():
        docs.append({"text": f"Module : {label}.", "source": f"Module — {label}",
                     "page": key, "kind": "module"})
    if os.path.exists(CATALOGUE):
        try:
            import pandas as pd
            df = pd.read_csv(CATALOGUE, encoding="utf-8-sig")
            for _, r in df.iterrows():
                txt = (f"Indicateur : {r.get('Indicateurs_FR', '')}. "
                       f"Secteur : {r.get('Secteur', '')}. "
                       f"Thème OCDE : {r.get('Thème OCDE', '')}. "
                       f"Source : {r.get('source', '')}. "
                       f"Période : {r.get('Date_debut', '')}–{r.get('Date_fin', '')}.")
                docs.append({"text": txt, "source": "Catalogue d'indicateurs", "page": "m4",
                             "kind": "indic",
                             "theme": str(r.get("Thème OCDE", "")).strip(),
                             "secteur": str(r.get("Secteur", "")).strip(),
                             "src": str(r.get("source", "")).strip()})
        except Exception:
            pass
    if os.path.isdir(DOCS_DIR):
        for path in sorted(glob.glob(os.path.join(DOCS_DIR, "**", "*"), recursive=True)):
            if not os.path.isfile(path):
                continue
            name = os.path.basename(path)
            for ch in _chunks(_read_document(path)):
                docs.append({"text": ch, "source": f"Document — {name}", "page": "m4",
                             "kind": "doc"})

    texts = [d["text"] for d in docs]
    try:
        from sklearn.feature_extraction.text import TfidfVectorizer
        vec = TfidfVectorizer(lowercase=True, ngram_range=(1, config.RAG_NGRAM_MAX), min_df=1)
        mat = vec.fit_transform(texts) if texts else None
        return {"docs": docs, "vec": vec, "mat": mat, "mode": "tfidf"}
    except Exception:
        return {"docs": docs, "vec": None, "mat": None, "mode": "keyword"}


def get_index():
    return build_index(tuple(_mtime(p) for p in _index_paths()))


@st.cache_data(show_spinner=False)
def _facets(sig):
    import pandas as pd
    df = pd.read_csv(CATALOGUE, encoding="utf-8-sig")
    def uniq(col):
        if col not in df.columns:
            return []
        return sorted({str(x).strip() for x in df[col].dropna()
                       if str(x).strip() and str(x).strip().lower() not in ("nan", "none")})
    return uniq("Thème OCDE"), uniq("Secteur"), uniq("source")


def facets():
    return _facets(_mtime(CATALOGUE))


# ----------------------------------------------------------------- récupération
def _match(d, filt):
    if not filt or d.get("kind") != "indic":
        return True
    for key in ("theme", "secteur", "src"):
        want = filt.get(key)
        if want and d.get(key) != want:
            return False
    return True


def retrieve(idx, query, k=config.RAG_TOP_K, filt=None):
    docs = idx["docs"]
    if not docs:
        return []
    if idx["mode"] == "tfidf" and idx["mat"] is not None:
        from sklearn.metrics.pairwise import cosine_similarity
        sims = cosine_similarity(idx["vec"].transform([query]), idx["mat"])[0]
        res = []
        for i in sims.argsort()[::-1]:
            if sims[i] <= 0:
                break
            if _match(docs[i], filt):
                res.append((docs[i], float(sims[i])))
            if len(res) >= k:
                break
        return res
    ql = set(re.findall(r"\w+", query.lower()))
    scored = []
    for d in docs:
        if not _match(d, filt):
            continue
        s = len(ql & set(re.findall(r"\w+", d["text"].lower())))
        if s:
            scored.append((d, s))
    scored.sort(key=lambda x: -x[1])
    return scored[:k]


# ----------------------------------------------------------------- génération
def _client():
    key = None
    try:
        key = st.secrets.get("ANTHROPIC_API_KEY")
    except Exception:
        key = None
    key = key or os.environ.get("ANTHROPIC_API_KEY")
    if not key:
        return None
    try:
        import anthropic
        return anthropic.Anthropic(api_key=key)
    except Exception:
        return None


SYS_BASE = ("Tu es l'assistant de la plateforme VizionAyiti2030 (Bureau du Coordonnateur "
            "Résident des Nations Unies — Haïti). Réponds de façon concise et factuelle, "
            "UNIQUEMENT à partir du CONTEXTE fourni. Si l'information n'y figure pas, "
            "dis-le clairement. Cite tes sources avec leur numéro [n]. Termine en orientant "
            "l'utilisateur vers le module pertinent et son lien interne (ex. ?page=m1).")


def _sys(lang):
    return SYS_BASE + (" Answer in English." if lang == "English" else " Réponds en français.")


def answer(idx, query, history, lang="Français", model=None, filt=None):
    hits = retrieve(idx, query, filt=filt)
    context = "\n\n".join(f"[{i+1}] ({h['source']}) {h['text']}"
                          for i, (h, _) in enumerate(hits)) or "(aucun extrait pertinent)"
    client = _client()
    if client is None:
        extraits = "\n\n".join(f"• ({h['source']}) {h['text'][:300]}" for h, _ in hits[:4])
        return ("⚠️ Clé API Anthropic absente (ANTHROPIC_API_KEY). Extraits les plus "
                "pertinents du magasin de données :\n\n" + (extraits or "—")), hits
    msgs = [{"role": r, "content": c} for r, c in history[-config.LLM_HISTORY_TURNS:]]
    msgs.append({"role": "user", "content": f"CONTEXTE :\n{context}\n\nQUESTION : {query}"})
    try:
        resp = client.messages.create(model=(model or config.LLM_MODEL),
                                      max_tokens=config.LLM_MAX_TOKENS,
                                      temperature=config.LLM_TEMPERATURE,
                                      system=_sys(lang), messages=msgs)
        return resp.content[0].text, hits
    except Exception as e:
        return f"Erreur lors de l'appel au modèle : {e}", hits


# ----------------------------------------------------------------- flux rapport
def _report_intent(q):
    return bool(re.search(r"rapport|report|docx|\bword\b|génér|gener|produi|produce|"
                          r"télécharg|download", q.lower()))


def _make_report(ss, lab, ind, lang):
    try:
        data, name = rapport.build_report(lab, ind, lang)
        ss["asst_report"] = (name, data)
        ss["asst_msgs"].append(("assistant", f"✅ Rapport généré pour « {lab} » ({lang}). "
                                "Téléchargez-le ci-dessous."))
    except Exception as e:
        ss["asst_msgs"].append(("assistant", f"Erreur lors de la génération : {e}"))
    ss["asst_pending"] = None


def _handle_query(ss, q):
    q = (q or "").strip()
    if not q:
        return
    ss["asst_answer_audio"] = None
    ss["asst_msgs"].append(("user", q))
    if _report_intent(q):
        hit = rapport.find_indicator(q)
        if hit:
            ss["asst_pending"] = hit
            ss["asst_msgs"].append(("assistant",
                f"Je peux générer un rapport d'analyse sur « {hit[0]} ». "
                "**Dans quelle langue ?** (choisissez ci-dessous)"))
        else:
            ss["asst_msgs"].append(("assistant",
                "Pour quel indicateur souhaitez-vous le rapport ? Précisez son nom, "
                "ou utilisez l'onglet **Analyse** (?page=analyse)."))
        return
    lang = ss.get("asst_lang", "Français")
    with st.spinner("Recherche dans le magasin de données…"):
        ans, hits = answer(get_index(), q, ss["asst_msgs"], lang,
                           ss.get("asst_model"), ss.get("asst_filt"))
    srcs = sorted({h["source"] for h, _ in hits})
    if srcs:
        ans += "\n\n— *Sources : " + " · ".join(srcs[:5]) + "*"
    plain = re.sub(r"[*_#\[\]]", "", ans.split("\n\n— *Sources")[0])
    creole = None
    want_voice = ss.get("asst_tts")
    if (lang == "Kreyòl" or want_voice) and translate.available():
        try:
            with st.spinner("Tradiksyon an kreyòl…"):
                creole = translate.to_creole(plain)
        except Exception:
            creole = None
    if lang == "Kreyòl" and creole:
        ans += "\n\n🇭🇹 **Kreyòl :** " + creole
    ss["asst_msgs"].append(("assistant", ans))
    if want_voice:
        au = voice.synthesize((creole or plain)[:config.TTS_MAX_CHARS])
        ss["asst_answer_audio"] = au if au else None


def _new_chat(ss):
    for k in ("asst_msgs", "asst_report", "asst_answer_audio", "asst_pending"):
        ss[k] = [] if k == "asst_msgs" else None


# ----------------------------------------------------------------- widget
def render_widget():
    ss = st.session_state
    ss.setdefault("asst_open", False)
    ss.setdefault("asst_msgs", [])
    ss.setdefault("asst_pending", None)
    ss.setdefault("asst_report", None)
    ss.setdefault("asst_answer_audio", None)
    ss.setdefault("asst_mic_hash", None)
    ss.setdefault("asst_lang", "Français")
    ss.setdefault("asst_model", config.LLM_MODEL)
    ss.setdefault("asst_filt", None)

    try:
        from streamlit_float import float_init
        float_init()
        has_float = True
    except Exception:
        has_float = False

    st.markdown("""
    <style>
      .asst-dock [data-testid="stForm"]{border:none;padding:0;}
      .asst-dock .stButton button{border-radius:9px;}
      .asst-dock [data-testid="stExpander"]{border:1px solid #e6e8ee;border-radius:9px;}
    </style>""", unsafe_allow_html=True)

    cont = st.container()
    with cont:
        st.markdown('<div class="asst-dock"></div>', unsafe_allow_html=True)
        if not ss["asst_open"]:
            if st.button("💬 Assistant", key="asst_open_btn", use_container_width=True):
                ss["asst_open"] = True
                st.rerun()
        else:
            h1, h2, h3 = st.columns([3, 1.4, 0.8])
            h1.markdown("**🛈 Assistant VizionAyiti2030**")
            if h2.button("✏️ Nouveau", key="asst_new"):
                _new_chat(ss)
                st.rerun()
            if h3.button("✕", key="asst_close_btn"):
                ss["asst_open"] = False
                st.rerun()

            # menu repliable : modèle, langue, filtres
            with st.expander("⚙️ Filtres & réglages"):
                ss["asst_model"] = st.selectbox("Modèle", MODELS,
                    index=MODELS.index(ss["asst_model"]) if ss["asst_model"] in MODELS else 0,
                    key="asst_model_sb")
                ss["asst_lang"] = st.radio("Langue des réponses", LANGS,
                    index=LANGS.index(ss["asst_lang"]), horizontal=True, key="asst_lang_rb")
                themes, secteurs, sources = facets()
                th = st.selectbox("Thème OCDE", ["Tous"] + themes, key="asst_f_theme")
                se = st.selectbox("Secteur", ["Tous"] + secteurs, key="asst_f_sec")
                so = st.selectbox("Source", ["Tous"] + sources, key="asst_f_src")
                ss["asst_filt"] = {"theme": None if th == "Tous" else th,
                                   "secteur": None if se == "Tous" else se,
                                   "src": None if so == "Tous" else so}

            for role, content in ss["asst_msgs"][-8:]:
                with st.chat_message("user" if role == "user" else "assistant"):
                    st.markdown(content)

            # questions suggérées (chat vide)
            if not ss["asst_msgs"]:
                st.caption("Questions suggérées :")
                for i, qs in enumerate(SUGGESTIONS):
                    if st.button(qs, key=f"asst_sug_{i}", use_container_width=True):
                        _handle_query(ss, qs)
                        st.rerun()

            if ss["asst_answer_audio"]:
                st.audio(ss["asst_answer_audio"])

            if ss["asst_report"]:
                name, data = ss["asst_report"]
                st.download_button("⬇️ " + name, data, file_name=name, key="asst_dl",
                    mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
                    use_container_width=True)

            if ss["asst_pending"]:
                lab, ind = ss["asst_pending"]
                b1, b2 = st.columns(2)
                if b1.button("🇫🇷 Français", key="asst_fr", use_container_width=True):
                    _make_report(ss, lab, ind, "Français")
                    st.rerun()
                if b2.button("🇬🇧 English", key="asst_en", use_container_width=True):
                    _make_report(ss, lab, ind, "English")
                    st.rerun()

            if voice.available():
                ss["asst_tts"] = st.checkbox("🔊 Lire la réponse (kreyòl)",
                                             value=ss.get("asst_tts", False), key="asst_tts_cb")
                try:
                    audio = st.audio_input("🎙️ Parler (kreyòl)", key="asst_mic")
                except Exception:
                    audio = None
                if audio is not None:
                    b = audio.getvalue()
                    hsh = hashlib.md5(b).hexdigest()
                    if hsh != ss["asst_mic_hash"]:
                        ss["asst_mic_hash"] = hsh
                        with st.spinner("Transcription…"):
                            text, err = voice.transcribe(b)
                        if err:
                            ss["asst_msgs"].append(("assistant", "⚠️ " + err))
                        elif text:
                            _handle_query(ss, text)
                        st.rerun()

            with st.form("asst_form", clear_on_submit=True):
                q = st.text_input("q", label_visibility="collapsed",
                                  placeholder="Posez votre question…")
                sent = st.form_submit_button("Envoyer", use_container_width=True)
            if sent and q.strip():
                _handle_query(ss, q)
                st.rerun()

    if has_float:
        try:
            cont.float(
                "bottom:1.2rem; right:1.2rem; width:430px; max-height:84vh; "
                "overflow-y:auto; background:#fff; border:1px solid #e6e8ee; "
                "border-radius:14px; box-shadow:0 12px 34px rgba(20,34,77,.22); "
                "padding:14px 16px; z-index:99999;")
        except Exception:
            pass