"""Vector store utilities using ChromaDB and Ollama embeddings.

This module provides functions to create a persistent Chroma vector store
and to load documents from a directory into the store.
"""

from pathlib import Path
from typing import Iterable

from langchain_chroma import Chroma
from langchain_ollama import OllamaEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

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

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

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

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

# ---------------------------------------------------------------------------
# Load documents from a directory into the vector store.
# ---------------------------------------------------------------------------

def load_documents(directory: str, vectorstore: Chroma) -> None:
    """Load all .txt and .md files from *directory* into *vectorstore*.

    The documents are split into chunks using a RecursiveCharacterTextSplitter
    before being added to the vector store.
    """
    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    docs = []
    for path in Path(directory).rglob("*.txt"):
        docs.append(path.read_text(encoding="utf-8"))
    for path in Path(directory).rglob("*.md"):
        docs.append(path.read_text(encoding="utf-8"))
    if not docs:
        return
    # Split the documents into chunks.
    chunks = splitter.split_text("\n\n".join(docs))
    # Create LangChain Document objects.
    from langchain.docstore.document import Document
    documents = [Document(page_content=chunk) for chunk in chunks]
    vectorstore.add_documents(documents)

# ---------------------------------------------------------------------------
# End of vectorstore.py
# ---------------------------------------------------------------------------