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

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

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

# ---------------------------------------------------------------------------
# Vector store creation
# ---------------------------------------------------------------------------

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 files are stored.

    Returns
    -------
    Chroma
        A Chroma vector store instance.
    """
    # Ensure directory exists
    Path(persist_directory).mkdir(parents=True, exist_ok=True)
    # Use Ollama embeddings
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
    # Create Chroma store
    vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
    return vectorstore

# ---------------------------------------------------------------------------
# Document loading
# ---------------------------------------------------------------------------

def _load_text_files(directory: str) -> List[str]:
    """Load all .txt and .md files from a directory into a list of strings."""
    texts = []
    for root, _, files in os.walk(directory):
        for file in files:
            if file.lower().endswith(('.txt', '.md')):
                path = Path(root) / file
                try:
                    content = path.read_text(encoding="utf-8")
                    texts.append(content)
                except Exception as e:
                    print(f"Failed to read {path}: {e}")
    return texts


def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200) -> None:
    """Load documents from a directory into the provided vector store.

    Parameters
    ----------
    directory: str
        Path to the directory containing .txt/.md files.
    vectorstore: Chroma
        The vector store to add documents to.
    chunk_size: int, optional
        Maximum size of each chunk.
    chunk_overlap: int, optional
        Number of characters to overlap between chunks.
    """
    texts = _load_text_files(directory)
    if not texts:
        print("No text files found in the directory.")
        return

    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    docs = []
    for text in texts:
        docs.extend(splitter.split_text(text))

    # Add documents to Chroma
    vectorstore.add_texts(docs)
    print(f"Loaded {len(docs)} chunks into the vector store.")

# ---------------------------------------------------------------------------
# Example usage (uncomment to run directly)
# ---------------------------------------------------------------------------
# if __name__ == "__main__":
#     store = create_vectorstore()
#     load_documents("documents", store)
""