"""Vector store utilities for ChromaDB.

This module provides functions to create a persistent ChromaDB vector store using
Ollama embeddings and to load documents from a directory into the store.
"""

from pathlib import Path

from langchain_chroma import Chroma
from langchain_ollama import OllamaEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document


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

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

    Returns
    -------
    Chroma
        A 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) -> None:
    """Load .txt and .md files from *directory* into *vectorstore*.

    The documents are split into chunks using ``RecursiveCharacterTextSplitter``
    before being added to the vector store.
    """
    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    docs = []
    for file_path in Path(directory).glob("*"):
        if file_path.suffix.lower() not in {".txt", ".md"}:
            continue
        with open(file_path, "r", encoding="utf-8") as f:
            text = f.read()
        chunks = splitter.split_text(text)
        docs.extend([Document(page_content=c, metadata={"source": str(file_path)}) for c in chunks])
    if docs:
        vectorstore.add_documents(docs)
        print(f"Loaded {len(docs)} chunks from {directory} into ChromaDB.")
    else:
        print(f"No .txt/.md files found in {directory}.")