"""
GENAURA - Script d'indexation des documents La Poste dans ChromaDB
Exécuter une seule fois avant de lancer l'application.
Prérequis : Ollama doit être lancé avec le modèle nomic-embed-text.
  ollama pull nomic-embed-text
"""

import os
import chromadb
from chromadb.utils import embedding_functions


def main():
    # Initialisation ChromaDB (stockage local persistant)
    client = chromadb.PersistentClient(path="./chroma_db")

    # Fonction d'embedding via Ollama (gratuit, local)
    ollama_ef = embedding_functions.OllamaEmbeddingFunction(
        url="http://localhost:11434/api/embeddings",
        model_name="nomic-embed-text"
    )

    # Créer ou récupérer la collection
    collection = client.get_or_create_collection(
        name="docs_la_poste",
        embedding_function=ollama_ef
    )

    # Charger et indexer les documents
    docs_dir = "./documents_la_poste"
    count = 0

    for filename in os.listdir(docs_dir):
        if filename.endswith(".md"):
            filepath = os.path.join(docs_dir, filename)
            with open(filepath, "r", encoding="utf-8") as f:
                content = f.read()

            # Ajouter ou mettre à jour le document
            collection.upsert(
                documents=[content],
                ids=[filename],
                metadatas=[{"source": filename}]
            )
            count += 1
            print(f"  Indexé : {filename}")

    print(f"\n{count} documents indexés dans ChromaDB.")
    print("Base vectorielle prête dans ./chroma_db/")


if __name__ == "__main__":
    print("=== GENAURA - Indexation des documents La Poste ===\n")
    main()