"""Utilities for creating and populating a Chroma vector store.

This module provides two helper functions:

* ``create_vectorstore`` – creates a Chroma collection backed by Ollama embeddings.
* ``load_documents`` – reads ``.txt``/``.md`` files, splits them into chunks and adds them to the collection.

The vector store is persisted in ``./chroma_db`` by default.
"""

from pathlib import Path
from typing import List

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

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

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

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

    Parameters
    ----------
    persist_directory: str
        Directory where the vector store will be persisted.

    Returns
    -------
    Chroma
        A Chroma collection ready for adding documents.
    """
    embeddings = OllamaEmbeddings(model=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*, split them into chunks and add to *vectorstore*.

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

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

    if docs:
        vectorstore.add_texts(docs)
    else:
        print("[vectorstore] No documents found in", directory)

# ---------------------------------------------------------------------------
# Example usage (uncomment to run as a script)
# ---------------------------------------------------------------------------
# if __name__ == "__main__":
#     store = create_vectorstore()
#     load_documents("documents", store)
#     print("Vector store populated.")
"""