import os
from pathlib import Path
from typing import List

from langchain_ollama import OllamaEmbeddings
from langchain_chroma import Chroma
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 backed by Ollama embeddings.

    Parameters
    ----------
    persist_directory: str
        Directory where Chroma will persist its data.

    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, chunk_size: int = 1000, chunk_overlap: int = 200) -> None:
    """Load all .txt and .md files from *directory* into *vectorstore*.

    The function reads files, splits them into chunks using
    ``RecursiveCharacterTextSplitter`` and adds the resulting
    :class:`~langchain.docstore.document.Document` objects to the
    vector store.
    """
    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    docs: List[Document] = []
    for path in Path(directory).glob("**/*"):
        if path.suffix.lower() not in {".txt", ".md"}:
            continue
        text = path.read_text(encoding="utf-8")
        docs.extend(splitter.split_text(text))
    # Convert list of strings to Document objects
    documents = [Document(page_content=chunk, metadata={"source": str(p)}) for chunk in docs]
    vectorstore.add_documents(documents)

# If this module is executed directly, load the default documents folder.
if __name__ == "__main__":
    vs = create_vectorstore()
    load_documents("documents", vs)
    print("Vector store populated.")