"""Vector store utilities for the FAQ bot.

This module handles loading markdown files into a persistent Chroma vector store
and provides a tool that performs semantic search over the stored documents.
"""

import os
from pathlib import Path
from typing import List

from langchain_ollama import OllamaEmbeddings
from langchain_chroma import Chroma
from langchain_core.documents import Document

from config import DATA_DIR, CHROMA_PERSIST_DIR, CHROMA_TOP_K
from chunker import chunk_text

# Initialize embeddings once
embeddings = OllamaEmbeddings(model="nomic-embed-text")

# Create or load the Chroma collection
vector_store = Chroma(
    collection_name="faq_collection",
    embedding_function=embeddings,
    persist_directory=CHROMA_PERSIST_DIR,
)

def load_faq_to_chroma() -> None:
    """Load all markdown files from DATA_DIR into the Chroma vector store.

    The function reads each .md file, splits it into chunks, creates Document
    objects, and adds them to the vector store. The store is persisted to disk.
    """
    docs: List[Document] = []
    for md_file in Path(DATA_DIR).glob("*.md"):
        text = md_file.read_text(encoding="utf-8")
        chunks = chunk_text(text)
        for i, chunk in enumerate(chunks):
            doc = Document(page_content=chunk, metadata={"source": md_file.name, "chunk_id": i})
            docs.append(doc)
    if docs:
        vector_store.add_documents(docs)
        vector_store.persist()

def search_course_docs(query: str, k: int = CHROMA_TOP_K) -> str:
    """Perform a semantic search over the FAQ collection.

    Parameters
    ----------
    query: str
        The user question.
    k: int, optional
        Number of top results to return.

    Returns
    -------
    str
        Concatenated top-k documents' page content.
    """
    results = vector_store.similarity_search_with_score(query, k=k)
    # results is list of (Document, score)
    snippets = [doc.page_content for doc, _ in results]
    return "\n\n---\n\n".join(snippets)