"""
Vector store utilities for ChromaDB with Ollama embeddings.

Functions:
- create_vectorstore(persist_directory="./chroma_db") -> chromadb.Chroma
- load_documents(directory: str, vectorstore) -> None
"""
import os
from pathlib import Path
from typing import List

from langchain_ollama import OllamaEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_chroma import Chroma
from langchain.docstore.document import Document


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

    Parameters
    ----------
    persist_directory : str, optional
        Directory where the Chroma database is stored. Defaults to ``./chroma_db``.

    Returns
    -------
    chromadb.Chroma
        A Chroma instance backed by Ollama embeddings.
    """
    # Ensure directory exists
    Path(persist_directory).mkdir(parents=True, exist_ok=True)
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
    return Chroma(persist_directory=persist_directory, embedding_function=embeddings)


def _load_text_files(directory: str) -> List[Document]:
    """Load all .txt and .md files from *directory* into Documents.

    Parameters
    ----------
    directory : str
        Path to the folder containing documents.

    Returns
    -------
    list[langchain.docstore.document.Document]
        List of Document objects with ``page_content`` set to file text and
        ``metadata`` containing the source path.
    """
    docs: List[Document] = []
    for root, _, files in os.walk(directory):
        for fname in files:
            if not fname.lower().endswith(('.txt', '.md')):
                continue
            full_path = Path(root) / fname
            text = full_path.read_text(encoding="utf-8")
            docs.append(Document(page_content=text, metadata={"source": str(full_path)}))
    return docs


def load_documents(directory: str, vectorstore: Chroma) -> None:
    """Chunk documents from *directory* and add them to *vectorstore*.

    Parameters
    ----------
    directory : str
        Folder with .txt/.md files.
    vectorstore : chromadb.Chroma
        The vector store instance returned by ``create_vectorstore``.
    """
    docs = _load_text_files(directory)
    if not docs:
        return
    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    chunks = splitter.split_documents(docs)
    vectorstore.add_documents(chunks)

# Example usage (uncomment for manual testing):
# if __name__ == "__main__":
#     vs = create_vectorstore()
#     load_documents("documents", vs)
"""