"""
Vector store implementation using ChromaDB.

Provides methods to add documents and perform similarity search.
"""
import os
from pathlib import Path
from typing import List, Dict

from chromadb import Client as ChromadbClient
from chromadb.config import Settings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_ollama import OllamaEmbeddings

# Ensure persistent directory exists
CHROMA_DIR = Path("./chroma_db")
CHROMA_DIR.mkdir(parents=True, exist_ok=True)

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

# Chroma client with persistence
client = ChromadbClient(Settings(persist_directory=str(CHROMA_DIR)))
collection_name = "knowledge"
if collection_name not in client.list_collections():
    client.create_collection(name=collection_name)
col = client.get_or_create_collection(name=collection_name)

# Chunker from chunker.py
from chunker import chunker

class VectorStore:
    def add_documents(self, docs: List[str], metadatas: List[Dict]):
        """Add documents to the collection.

        Parameters
        ----------
        docs: list of raw text strings.
        metadatas: list of metadata dicts corresponding to each doc.
        """
        # Split into chunks and embed
        all_chunks = []
        all_metadatas = []
        for doc, meta in zip(docs, metadatas):
            chunks = chunker.split_text(doc)
            all_chunks.extend(chunks)
            all_metadatas.extend([meta] * len(chunks))

        # Generate embeddings via Ollama
        embeds = embeddings.embed_documents(all_chunks)
        ids = [f"chunk_{i}" for i in range(len(all_chunks))]
        col.add(ids=ids, documents=all_chunks, metadatas=all_metadatas, embeddings=embeds)

    def similarity_search(self, query: str, k: int = 5) -> List[Dict]:
        """Return top-k similar chunks with metadata."""
        results = col.query(
            query_texts=[query],
            n_results=k,
            include=['documents', 'metadatas'],
        )
        # results is dict with keys documents, metadatas
        docs = results["documents"][0]
        metas = results["metadatas"][0]
        return [{"content": d, "metadata": m} for d, m in zip(docs, metas)]