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

import os
from pathlib import Path

from langchain_chroma import Chroma
from langchain_ollama import OllamaEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader, UnstructuredMarkdownLoader

# ---------------------------------------------------------------------------
# Create a persistent Chroma vector store
# ---------------------------------------------------------------------------

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

    Parameters
    ----------
    persist_directory: str
        Directory where the Chroma DB will be persisted.

    Returns
    -------
    Chroma
        A Chroma vector store instance.
    """
    os.makedirs(persist_directory, exist_ok=True)
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
    return Chroma(persist_directory=persist_directory, embedding_function=embeddings)

# ---------------------------------------------------------------------------
# Load documents from a directory and add them to the vector store
# ---------------------------------------------------------------------------

def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 500, chunk_overlap: int = 50) -> None:
    """Load `.txt` and `.md` files from *directory*, split them into chunks, and add to *vectorstore*.

    Parameters
    ----------
    directory: str
        Directory containing the source documents.
    vectorstore: Chroma
        The vector store to which documents will be added.
    chunk_size: int, optional
        Maximum chunk size in characters.
    chunk_overlap: int, optional
        Number of overlapping characters between consecutive chunks.
    """
    loader_classes = {
        ".txt": TextLoader,
        ".md": UnstructuredMarkdownLoader,
    }
    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)

    docs = []
    for root, _, files in os.walk(directory):
        for file in files:
            ext = Path(file).suffix.lower()
            if ext not in loader_classes:
                continue
            loader = loader_classes[ext](os.path.join(root, file))
            loaded_docs = loader.load()
            docs.extend(loaded_docs)

    if not docs:
        return

    # Split documents into smaller chunks
    split_docs = splitter.split_documents(docs)
    vectorstore.add_documents(split_docs)

# ---------------------------------------------------------------------------
# Example usage
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    # This block is only executed when running the module directly.
    store = create_vectorstore()
    load_documents("documents", store)
    print("Vector store populated.")