#!/usr/bin/env python3
"""
Indexation incrémentale des fiches avec BGE‑M3 (FlagEmbedding) + FAISS
---------------------------------------------------------------------
• Parcourt récursivement le dossier Fiches (ext .md .MD .markdown .txt)
• Découpe chaque fichier en blocs de ~800 tokens (chevauchement 100)
• Encode uniquement les passages provenant de fichiers **nouveaux ou modifiés**
  depuis la dernière indexation (basé sur l'horodatage mtime).
• Les embeddings sont ajoutés à l'index existant sans toucher aux anciens.
• Écrit/Met à jour : corpus.idx (vecteurs) + corpus.meta.json (métadonnées)

➡ Robustesse améliorée : le script détecte tous les formats de retour possibles
   de `BGEM3FlagModel.encode` (ndarray ou dict avec clef `embedding`, etc.).
"""

from pathlib import Path
import json
import re
import time
from datetime import datetime
import faiss
import numpy as np
from FlagEmbedding import BGEM3FlagModel

# --- Paramètres -------------------------------------------------------------
ROOT = Path("Fiches")                       # répertoire local des fiches
MODEL_NAME = "BAAI/bge-m3"                  # embedder multilingue, MIT
CHUNK = 800                                  # taille cible d'un bloc (≈600 mots)
OVERLAP = 100                                # chevauchement pour la cohérence
INDEX_FILE = "corpus.idx"
META_FILE = "corpus.meta.json"
EXTENSIONS = ["*.md", "*.MD", "*.markdown", "*.txt"]
BATCH = 256                                   # plus grand batch : encode plus vite

# --- Fonctions utilitaires --------------------------------------------------

def split(text: str, chunk_size: int = CHUNK, overlap: int = OVERLAP):
    """Découpe un texte en morceaux de chunk_size mots avec overlap mots."""
    sentences = re.split(r"(?<=[\.!?])\s+", text)
    chunks, buf = [], []
    for s in sentences:
        buf.append(s)
        if len(" ".join(buf)) > chunk_size:
            chunks.append(" ".join(buf))
            buf = buf[-overlap:]
    if buf:
        chunks.append(" ".join(buf))
    return chunks


def gather_files(root: Path):
    for pattern in EXTENSIONS:
        yield from root.rglob(pattern)

# --- Chargement éventuel de l'index existant -------------------------------

def load_existing():
    if not Path(INDEX_FILE).exists():
        return None, [], 0  # pas d'index, pas de meta

    index = faiss.read_index(INDEX_FILE)
    meta = json.load(open(META_FILE, encoding="utf-8"))
    return index, meta, len(meta)

# --- Pipeline principal -----------------------------------------------------

def main():
    index, meta, existing = load_existing()

    # mapping chemin relatif ➜ mtime stocké
    meta_mtime = {m["file"]: m.get("mtime", 0) for m in meta}

    new_docs, new_meta = [], []
    files_scanned, files_updated = 0, 0

    for fp in gather_files(ROOT):
        files_scanned += 1
        rel = fp.relative_to(ROOT).as_posix()
        mtime = int(fp.stat().st_mtime)
        if meta_mtime.get(rel) == mtime:
            continue  # inchangé
        files_updated += 1
        text = fp.read_text(encoding="utf-8", errors="ignore")
        for i, chunk in enumerate(split(text)):
            new_docs.append(chunk)
            new_meta.append({"file": rel, "part": i, "mtime": mtime})

    if not new_docs:
        print("Aucun fichier nouveau ou modifié. Index à jour ✔︎")
        return

    print(
        f"Nouveaux passages : {len(new_docs)} issus de {files_updated} fiches ("\
        f"{files_scanned} fiches scannées). Génération des embeddings…"
    )

    model = BGEM3FlagModel(MODEL_NAME, device="cpu")
    emb_out = model.encode(new_docs, batch_size=BATCH)

    # --- Gestion robuste du format de sortie --------------------------------
    if isinstance(emb_out, np.ndarray):
        emb = emb_out
    elif isinstance(emb_out, dict):
        for key in ("embedding", "embeddings", "sentence_embeds", "sentence_embedding"):
            if key in emb_out:
                emb = np.asarray(emb_out[key])
                break
        else:
            emb = np.asarray(next(iter(emb_out.values())))
    else:  # liste de tableaux ?
        emb = np.asarray(emb_out)

    emb = emb.astype("float32")
    emb /= np.linalg.norm(emb, axis=1, keepdims=True) + 1e-12

    if index is None:
        index = faiss.IndexFlatIP(emb.shape[1])
    index.add(emb)

    # Mettre à jour les métadonnées (anciens + nouveaux)
    meta.extend(new_meta)
    faiss.write_index(index, INDEX_FILE)
    json.dump(meta, open(META_FILE, "w", encoding="utf-8"), ensure_ascii=False, indent=2)

    print(
        f"Index mis à jour : {len(meta)} vecteurs au total ("\
        f"+{len(new_docs)}). Dernière maj : {datetime.now().isoformat(timespec='seconds')}"
    )

if __name__ == "__main__":
    t0 = time.time()
    main()
    print(f"Terminé en {time.time() - t0:.1f} s")