"""ChromaDB vector store with Ollama embeddings and document chunking.""
import os
from typing import List

from langchain_chroma import Chroma
from langchain_ollama import OllamaEmbeddings
from langchain_core.documents import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter

COLLECTION_NAME = "knowledge_base"
PERSIST_DIR = os.getenv("CHROMA_PERSIST_DIR", "./chroma_db")

embeddings = OllamaEmbeddings(model="nomic-embed-text")

_store: Chroma | None = None


def get_vector_store() -> Chroma:
    """Return singleton ChromaDB instance (persisted to disk)."""
    global _store
    if _store is None:
        _store = Chroma(
            collection_name=COLLECTION_NAME,
            embedding_function=embeddings,
            persist_directory=PERSIST_DIR,
        )
    return _store


def add_documents(texts: List[str], chunk_size: int = 500, chunk_overlap: int = 100) -> int:
    """Split texts into chunks and index them in ChromaDB.

    Args:
        texts: list of raw text strings
        chunk_size: maximum chunk length in characters
        chunk_overlap: overlap between consecutive chunks

    Returns:
        number of chunks added to the store
    """
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
    )
    docs: List[Document] = splitter.create_documents(texts)
    if docs:
        get_vector_store().add_documents(docs)
    return len(docs)


def similarity_search(query: str, k: int = 5) -> List[Document]:
    """Semantic search in ChromaDB.

    Args:
        query: natural language search query
        k: number of top results to return

    Returns:
        list of the most relevant Document objects
    """
    return get_vector_store().similarity_search(query, k=k)