"""Vector store utilities using ChromaDB and Ollama embeddings.

This module provides functions to create a persistent ChromaDB vector store and
load documents from a directory into it. The documents are split into
manageable chunks using a recursive character text splitter.
"""

from pathlib import Path
from typing import Iterable

from langchain_chroma import Chroma
from langchain_ollama import OllamaEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

# Default persistence directory
DEFAULT_PERSIST_DIR = "./chroma_db"

# Default chunking parameters – these can be tuned
DEFAULT_CHUNK_SIZE = 1000
DEFAULT_CHUNK_OVERLAP = 200


def create_vectorstore(persist_directory: str = DEFAULT_PERSIST_DIR) -> Chroma:
    """Create or load a Chroma vector store.

    Parameters
    ----------
    persist_directory: str
        Directory where the ChromaDB files are stored.

    Returns
    -------
    Chroma
        An instance of the Chroma vector store backed by the given directory.
    """
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
    vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
    return vectorstore


def _load_text_files(directory: str | Path) -> Iterable[str]:
    """Yield the contents of all .txt and .md files in *directory*.

    Parameters
    ----------
    directory: str | Path
        Path to the directory containing the documents.
    """
    directory = Path(directory)
    for file_path in directory.rglob("*.txt"):
        yield file_path.read_text(encoding="utf-8")
    for file_path in directory.rglob("*.md"):
        yield file_path.read_text(encoding="utf-8")


def load_documents(directory: str | Path, vectorstore: Chroma) -> None:
    """Load documents from *directory* into *vectorstore*.

    The documents are split into chunks using a recursive character splitter
    and then added to the vector store. Existing documents are not removed –
    this function simply appends new data.

    Parameters
    ----------
    directory: str | Path
        Directory containing the source documents.
    vectorstore: Chroma
        The vector store to populate.
    """
    splitter = RecursiveCharacterTextSplitter(chunk_size=DEFAULT_CHUNK_SIZE, chunk_overlap=DEFAULT_CHUNK_OVERLAP)
    for text in _load_text_files(directory):
        chunks = splitter.split_text(text)
        vectorstore.add_texts(chunks)

    # Persist changes
    vectorstore.persist()

# End of vectorstore.py