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

import os
from pathlib import Path
from typing import List

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

# Create the vector store with persistence

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 persisted.

    Returns
    -------
    Chroma
        The Chroma vector store instance.
    """
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
    return Chroma(persist_directory=persist_directory, embedding_function=embeddings)

# Load documents from a directory and add them to the vector store

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

    Parameters
    ----------
    directory: str
        Path to the directory containing the documents.
    vectorstore: Chroma
        The vector store to which the documents will be added.
    """
    # Ensure the directory exists
    path = Path(directory)
    if not path.is_dir():
        raise FileNotFoundError(f"Directory {directory} does not exist")

    # Collect all .txt and .md files
    files = list(path.rglob("*.txt")) + list(path.rglob("*.md"))
    if not files:
        print(f"No .txt or .md files found in {directory}")
        return

    # Read and chunk the documents
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    documents = []
    for file_path in files:
        try:
            content = file_path.read_text(encoding="utf-8")
        except Exception as e:
            print(f"Failed to read {file_path}: {e}")
            continue
        # Split into chunks
        chunks = text_splitter.split_text(content)
        for i, chunk in enumerate(chunks):
            documents.append({
                "page_content": chunk,
                "metadata": {"source": str(file_path), "chunk_index": i},
            })

    if documents:
        vectorstore.add_documents(documents)
        print(f"Added {len(documents)} chunks to the vector store from {directory}")
    else:
        print("No documents were processed.")

# Example usage:
# if __name__ == "__main__":
#     store = create_vectorstore()
#     load_documents("./documents", store)