"""
Vector store utilities for the RAG agent.

Provides functions to create a ChromaDB vector store backed by Ollama embeddings
and to load documents from a directory into the store.
"""

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

# ---------------------------------------------------------------------------
# Configuration constants
# ---------------------------------------------------------------------------
DEFAULT_EMBEDDING_MODEL = "nomic-embed-text"
DEFAULT_PERSIST_DIR = "./chroma_db"

# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

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

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

    Returns
    -------
    Chroma
        An instance of the Chroma vector store.
    """
    embeddings = OllamaEmbeddings(model=DEFAULT_EMBEDDING_MODEL)
    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 all .txt and .md files from *directory*, chunk them and add to *vectorstore*.

    The function is idempotent – if the same files are loaded again, duplicates will
    not be created because Chroma will deduplicate based on the content hash.

    Parameters
    ----------
    directory: str
        Path to the folder containing the documents.
    vectorstore: Chroma
        The vector store to populate.
    chunk_size: int, optional
        Maximum number of characters per chunk.
    chunk_overlap: int, optional
        Number of characters to overlap between consecutive chunks.
    """
    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    docs: List[Document] = []

    for file_path in Path(directory).glob("**/*"):
        if file_path.suffix.lower() not in {".txt", ".md"}:
            continue
        text = file_path.read_text(encoding="utf-8")
        docs.extend(splitter.split_text(text))

    # Convert list of strings to list of Documents
    documents = [Document(page_content=chunk) for chunk in docs]

    if documents:
        vectorstore.add_documents(documents)
        vectorstore.persist()

# ---------------------------------------------------------------------------
# End of module
# ---------------------------------------------------------------------------