"""Utilities for creating and populating a ChromaDB vector store.

This module provides two functions:

- :func:`create_vectorstore` – creates a Chroma vector store backed by a local directory.
- :func:`load_documents` – reads all ``.txt`` and ``.md`` files from a directory, splits them into chunks using
  :class:`langchain_text_splitters.RecursiveCharacterTextSplitter`, and upserts the chunks into the
  provided vector store.

The vector store is persistent across runs – the ``persist_directory`` argument defaults to
``"./chroma_db"``.
"""

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 embedding model used by the vector store
EMBEDDING_MODEL = "nomic-embed-text"


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

    Parameters
    ----------
    persist_directory: str
        Directory where the vector store will be persisted.

    Returns
    -------
    Chroma
        A Chroma vector store instance.
    """
    embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
    return Chroma(persist_directory=persist_directory, embedding_function=embeddings)


def _read_text_files(directory: str) -> Iterable[str]:
    """Yield the content of all ``.txt`` and ``.md`` files in *directory*.
    """
    path = Path(directory)
    for file_path in path.rglob("*.txt"):
        yield file_path.read_text(encoding="utf-8")
    for file_path in path.rglob("*.md"):
        yield file_path.read_text(encoding="utf-8")


def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200) -> None:
    """Load documents from *directory* into *vectorstore*.

    The documents are split into chunks using :class:`RecursiveCharacterTextSplitter` and then
    upserted into the vector store.
    """
    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    for text in _read_text_files(directory):
        chunks = splitter.split_text(text)
        vectorstore.add_texts(chunks)

# End of vectorstore.py