import os
import chromadb
from chromadb.utils import embedding_functions

# Configurar persistencia en la carpeta raíz del proyecto
PERSIST_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), ".chroma_db")

class VectorStore:
    """Implementación nativa de base de datos vectorial usando ChromaDB para búsquedas locales."""
    
    def __init__(self, collection_name: str = "letxipu_docs"):
        # Asegurar que el directorio de persistencia exista
        os.makedirs(PERSIST_DIR, exist_ok=True)
        
        self.client = chromadb.PersistentClient(path=PERSIST_DIR)
        
        # Usar un modelo ligero local multilingüe o estándar
        self.embedding_fn = embedding_functions.SentenceTransformerEmbeddingFunction(
            model_name="paraphrase-multilingual-MiniLM-L12-v2"  # Soporta español e inglés muy bien
        )
        
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            embedding_function=self.embedding_fn
        )
        
    def add_documents(self, documents: list[str], metadatas: list[dict], ids: list[str]):
        """Añade documentos (chunks) a la base vectorial."""
        if not documents:
            return
        self.collection.add(
            documents=documents,
            metadatas=metadatas,
            ids=ids
        )
        
    def search(self, query: str, n_results: int = 5, filter_dict: dict = None) -> dict:
        """Busca los fragmentos semánticamente más similares a la consulta."""
        results = self.collection.query(
            query_texts=[query],
            n_results=n_results,
            where=filter_dict
        )
        return results

    def clear(self):
        """Elimina todos los documentos de la colección actual."""
        all_ids = self.collection.get().get("ids", [])
        if all_ids:
            self.collection.delete(ids=all_ids)