"""
Vector store implementation using ChromaDB.

Provides functions to add documents and perform semantic search.
"""

import os
from typing import List, Dict
from langchain_ollama import OllamaEmbeddings
from chromadb import Client
from chromadb.config import Settings

# Initialize embeddings model (Ollama)
embeddings = OllamaEmbeddings(model="nomic-embed-text")

# ChromaDB client – in‑memory by default, persistent folder "chromadb"
CHROMA_DIR = os.path.join(os.getcwd(), "chromadb")
client = Client(Settings(chroma_db_impl="duckdb+parquet", persist_directory=CHROMA_DIR))
collection_name = "rag_collection"

# Ensure collection exists
if collection_name not in client.list_collections():
    client.create_collection(name=collection_name)
col = client.get_or_create_collection(name=collection_name)

class VectorStore:
    def __init__(self, collection):
        self.collection = collection

    def add_document(self, doc_id: str, text: str, metadata: Dict | None = None) -> None:
        """Add a single document to the collection.

        Parameters
        ----------
        doc_id: str
            Unique identifier for the document.
        text: str
            Raw text content.
        metadata: dict, optional
            Additional key/value pairs stored with the vector.
        """
        vec = embeddings.embed_query(text)
        self.collection.add(ids=[doc_id], documents=[text], metadatas=[metadata or {}])

    def search(self, query: str, k: int = 5) -> List[Dict]:
        """Semantic search over the collection.

        Returns a list of dicts with keys: id, document, score, metadata.
        """
        results = self.collection.query(
            query_texts=[query], n_results=k, include=['documents', 'distances', 'metadatas']
        )
        hits = []
        for i in range(len(results["ids"][0])):
            hit = {
                "id": results["ids"][0][i],
                "document": results["documents"][0][i],
                "score": 1 - results["distances"][0][i],  # distance to similarity
                "metadata": results["metadatas"][0][i],
            }
            hits.append(hit)
        return hits

# Singleton instance for easy import
vector_store = VectorStore(col)