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

This module provides functions to create a persistent Chroma vector store and to load documents
from a directory into the store. Documents are split into chunks using
`RecursiveCharacterTextSplitter`.
"""

import os
from pathlib import Path
from typing import List

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

# Default persistence directory
DEFAULT_PERSIST_DIR = "./chroma_db"


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

    Parameters
    ----------
    persist_directory: str
        Directory where the Chroma DB will be stored.

    Returns
    -------
    Chroma
        The Chroma vector store instance.
    """
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
    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* into *vectorstore*.

    The files are read, split into chunks with a recursive character splitter and
    added to the Chroma collection.
    """
    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    docs: List[str] = []
    for file_path in Path(directory).rglob("*.txt"):
        docs.append(file_path.read_text(encoding="utf-8"))
    for file_path in Path(directory).rglob("*.md"):
        docs.append(file_path.read_text(encoding="utf-8"))

    if not docs:
        return
    # Split documents into chunks
    texts = splitter.split_text("\n\n".join(docs))
    # Create a list of dicts with metadata (optional)
    metadatas = [{"source": "local"} for _ in texts]
    vectorstore.add_texts(texts, metadatas=metadatas)

    # Persist changes
    vectorstore.persist()

    print(f"Loaded {len(texts)} chunks into ChromaDB.")