"""
RAG vector store using ChromaDB and Ollama embeddings.
"""

from pathlib import Path
from typing import List

import chromadb
from langchain.embeddings.ollama import OllamaEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.schema.document import Document
from langchain.vectorstores import Chroma

CHROMA_DIR = "./chroma_db"
EMBED_MODEL = "nomic-embed-text"


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

    Parameters
    ----------
    persist_directory : str, optional
        Directory where the Chroma database is stored. If it does not exist,
        it will be created automatically by Chroma.
    """
    embeddings = OllamaEmbeddings(model=EMBED_MODEL)
    # Chroma can use a local directory for persistence
    client = chromadb.PersistentClient(path=persist_directory)
    collection = client.get_or_create_collection(name="documents", embedding_function=embeddings)
    return Chroma(collection=collection, embedding_function=embeddings)


def load_documents(directory: str, vectorstore: Chroma) -> None:
    """Load all .txt and .md files from *directory*, chunk them and add to the vector store.

    The function does not return anything; it mutates the provided collection.
    """
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    docs: List[Document] = []
    for path in Path(directory).rglob("*.txt"):
        content = path.read_text(encoding="utf-8")
        docs.extend(text_splitter.create_documents([content], metadata={"source": str(path)}))
    for path in Path(directory).rglob("*.md"):
        content = path.read_text(encoding="utf-8")
        docs.extend(text_splitter.create_documents([content], metadata={"source": str(path)}))

    if docs:
        # Chroma expects a list of documents via add_documents
        vectorstore.add_documents(docs)