"""
Vector store utilities for the RAG agent.

Provides:
- create_vectorstore(persist_directory)
- load_documents(directory, vectorstore)
"""
import os
from pathlib import Path
from typing import List

from langchain_ollama import OllamaEmbeddings
from langchain_chroma import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter


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

    Parameters
    ----------
    persist_directory: str
        Directory where the Chroma DB will be stored. If it does not exist, it is created.
    """
    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_documents(directory: str, vectorstore: Chroma) -> None:
    """Load all .txt/.md files from *directory* into the vector store.

    The function reads each file, splits it with a RecursiveCharacterTextSplitter and adds the chunks to the collection.
    Existing documents are overwritten – this is fine for an init script.
    """
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    docs: List[str] = []
    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:
        print(f"No documents found in {directory}")
        return

    # Split and add to vectorstore
    chunks = text_splitter.split_documents([{"content": d} for d in docs])
    vectorstore.add_texts([c["content"] for c in chunks])
    print(f"Loaded {len(chunks)} chunks into ChromaDB.")