"""
MODULE: VECTOR SEARCH ENGINE (OG-RAG FIXED)
===========================================
Correction : Ajout de la gestion des métadonnées (meta_map) pour éviter le KeyError.
"""
import functools
import os

import faiss
import numpy as np


@functools.lru_cache(maxsize=2)
def charger_modele(model_name):
    """Un seul modele en memoire par processus, partage par toutes les sessions."""
    from sentence_transformers import SentenceTransformer
    os.environ['HF_HUB_DISABLE_SYMLINKS_WARNING'] = '1'
    return SentenceTransformer(model_name)


class SemanticIndex:
    def __init__(self, model_name='all-MiniLM-L6-v2'):
        self.model = charger_modele(model_name)
        self.index = None
        self.uris = [] 
        self.content_map = {} 
        self.meta_map = {} # <--- AJOUT CRITIQUE

    def build_from_ontology_blocks(self, blocks):
        """Vectorise les Hyper-blocs enrichis"""
        print(f"[VECTOR] Indexation de {len(blocks)} Hyper-blocs...")
        
        corpus = [b['text'] for b in blocks]
        self.uris = [b['uri'] for b in blocks]
        self.content_map = {b['uri']: b['text'] for b in blocks}
        self.meta_map = {b['uri']: b['metadata'] for b in blocks} # <--- ON STOCKE LES METAS

        embeddings = self.model.encode(corpus, show_progress_bar=True)
        dimension = embeddings.shape[1]
        self.index = faiss.IndexFlatL2(dimension)
        self.index.add(np.array(embeddings).astype('float32'))
        
        print("[VECTOR] Index OG-RAG prêt.")

    def search(self, query, top_k=3):
        if not self.index: return []
        query_vec = self.model.encode([query]).astype('float32')
        dist, indices = self.index.search(query_vec, top_k)
        
        results = []
        for i, idx in enumerate(indices[0]):
            if idx != -1:
                uri = self.uris[idx]
                results.append({
                    "uri": uri,
                    "text": self.content_map.get(uri, ""),
                    "meta": self.meta_map.get(uri, {"type": "Unknown", "label": "?"}), # <--- ON RENVOIE LES METAS
                    "score": float(dist[0][i])
                })
        return results