"""
Vector store utilities for ChromaDB.
"""

from pathlib import Path

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


def create_vectorstore(persist_directory: str = "./chroma_db"):
    """Create or load a Chroma vector store.

    Parameters
    ----------
    persist_directory: str
        Directory where the Chroma database is persisted.

    Returns
    -------
    Chroma
        The Chroma vector store instance.
    """
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
    return Chroma(persist_directory=persist_directory, embedding_function=embeddings)


def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200):
    """Read all .txt and .md files from a directory, split them into chunks and add to the vectorstore.

    Parameters
    ----------
    directory: str
        Path to the folder containing documents.
    vectorstore: Chroma
        The vector store to which documents will be added.
    chunk_size: int
        Maximum number of characters per chunk.
    chunk_overlap: int
        Number of overlapping characters between consecutive chunks.
    """
    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    docs = []
    for file_path in Path(directory).rglob("*.txt"):
        docs.append(file_path.read_text(encoding="utf-8"))
    for file_path in Path(directory).rglob("*.md"):
        docs.append(file_path.read_text(encoding="utf-8"))
    if not docs:
        return
    texts = splitter.split_text("\n\n".join(docs))
    vectorstore.add_texts(texts)
    vectorstore.persist()