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

This module provides functions to create a persistent Chroma vector store and load
text documents from a directory into it.  The store is exposed via the global
``store`` variable so that other modules (e.g. tools) can access it.
"""

from pathlib import Path
from typing import List

from langchain_chroma import Chroma
from langchain_ollama import OllamaEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document

# Global store that will be initialised in ``create_vectorstore``.
store: Chroma | None = None

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 data will be persisted.
    """
    global store
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
    store = Chroma(
        collection_name="rag_collection",
        embedding_function=embeddings,
        persist_directory=persist_directory,
    )
    return store

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

    The documents are split using ``RecursiveCharacterTextSplitter`` before
    being added to the collection.
    """
    dir_path = Path(directory)
    txt_files = list(dir_path.rglob("*.txt")) + list(dir_path.rglob("*.md"))
    if not txt_files:
        print(f"No .txt or .md files found in {dir_path}")
        return

    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    docs: List[Document] = []
    for file_path in txt_files:
        text = file_path.read_text(encoding="utf-8")
        docs.extend(splitter.create_documents([text], metadata={"source": str(file_path)}))

    vectorstore.add_documents(docs)
    # Persist the collection to disk.
    vectorstore.persist()
    print(f"Loaded {len(docs)} documents from {dir_path} into Chroma.")

# Helper to get the global store.

def get_vectorstore() -> Chroma:
    if store is None:
        raise RuntimeError("Vector store has not been initialised. Call create_vectorstore() first.")
    return store

"""End of vectorstore.py"""