"""Module for creating and loading a Chroma vector store with Ollama embeddings."""

from pathlib import Path
from typing import List

from langchain_ollama import OllamaEmbeddings
from langchain_chroma import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.docstore.document 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
        The created Chroma vector store.
    """
    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) -> None:
    """Load .txt and .md files from a directory, split them into chunks, and add to the vector store.

    Parameters
    ----------
    directory: str
        Path to the directory containing documents.
    vectorstore: Chroma
        The vector store to add documents to.
    chunk_size: int
        Maximum size of each chunk.
    chunk_overlap: int
        Number of characters to overlap between chunks.
    """
    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    documents: List[Document] = []
    for file_path in Path(directory).glob("**/*"):
        if file_path.suffix.lower() in {".txt", ".md"}:
            text = file_path.read_text(encoding="utf-8")
            chunks = splitter.split_text(text)
            for chunk in chunks:
                documents.append(Document(page_content=chunk, metadata={"source": str(file_path)}))
    if documents:
        vectorstore.add_documents(documents)
        vectorstore.persist()