"""
Vector store utilities for ChromaDB with Ollama embeddings.
"""

import os
from pathlib import Path

from langchain_chroma import Chroma
from langchain_ollama import OllamaEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.docstore.document import Document

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

    Parameters
    ----------
    persist_directory: str
        Directory where the Chroma database is persisted.

    Returns
    -------
    Chroma
        A 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) -> None:
    """Load text and markdown files from *directory*, chunk them, and add to *vectorstore*.

    Parameters
    ----------
    directory: str
        Path to the folder containing .txt and .md files.
    vectorstore: Chroma
        The Chroma vector store to which documents will be added.
    """
    docs = []
    for file_path in Path(directory).rglob("*.*"):
        if file_path.suffix.lower() in ".txt .md".split():
            try:
                text = file_path.read_text(encoding="utf-8")
            except Exception as e:
                print(f"Failed to read {file_path}: {e}")
                continue
            docs.append(Document(page_content=text, metadata={"source": str(file_path)}))

    if not docs:
        print("No documents found to load.")
        return

    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    chunks = splitter.split_documents(docs)
    vectorstore.add_documents(chunks)
    print(f"Loaded {len(chunks)} chunks into the vector store.")