#!/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)

Temps gagné : pour 1 fiche ajoutée, seules ses quelques dizaines de passages
sont ré-encodés, pas les ~6 000 déjà en place ➜ ré‑indexation en quelques
secondes au lieu de minutes.
"""

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 = 128                                   # 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 = model.encode(new_docs, batch_size=BATCH, return_dict=False)
    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")