"""FAQ Bot with ChromaDB and a mock MCP tool.

The agent answers questions about course materials using a local Chroma vector store.
If the question is about course metadata (e.g., schedule), it calls a simple HTTP mock that returns JSON.  The agent is built with LangGraph.

Run with:
    python main.py

The script will load the FAQ files, build the vector store, and then enter an interactive loop.
"""

import json
import pathlib
from pathlib import Path

from langchain_ollama import ChatOllama
from langchain_chroma import Chroma
from langchain_text_splitters import MarkdownHeaderTextSplitter
from langchain_community.document_loaders import TextLoader
from langchain_community.embeddings import OllamaEmbeddings
from langgraph.prebuilt import create_react_agent
from langgraph.graph import StateGraph, MessagesState

# ---------- 1. Load FAQ files and create Chroma store ----------

def load_faq_to_chroma(data_dir: str = "data", persist_dir: str = "./chroma_faq") -> Chroma:
    """Read all .md files in data_dir, chunk them, and persist to Chroma."""
    Path(persist_dir).mkdir(parents=True, exist_ok=True)
    docs = []
    for md_file in Path(data_dir).glob("*.md"):
        loader = TextLoader(str(md_file), encoding="utf-8")
        docs.extend(loader.load())
    splitter = MarkdownHeaderTextSplitter(headers_to_split_on=["#", "##", "###"])
    split_docs = splitter.split_documents(docs)
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
    chroma = Chroma.from_documents(
        documents=split_docs,
        embedding=embeddings,
        persist_directory=persist_dir,
    )
    return chroma

# ---------- 2. MCP-style tool: fetch course metadata ----------

def fetch_course_meta(query: str) -> dict:
    """Mock MCP tool that returns course metadata.
    In production this would be an HTTP call to an MCP server.
    Here we simulate with a local JSON file.
    """
    meta_path = Path("mock_meta.json")
    if not meta_path.exists():
        meta = {"schedule": {"Monday": "10:00-12:00", "Wednesday": "14:00-16:00"}, "instructor": "Prof. Smith"}
        meta_path.write_text(json.dumps(meta, indent=2))
    else:
        meta = json.loads(meta_path.read_text())
    # Return the whole meta; the agent can filter as needed
    return meta

# ---------- 3. Agent setup ----------

def build_agent(chroma: Chroma):
    """Create a LangGraph agent that routes to either Chroma or the MCP tool."""
    def meta_tool(query: str):
        return json.dumps(fetch_course_meta(query), indent=2)
    tools = {"fetch_course_meta": meta_tool}
    agent = create_react_agent(ChatOllama(model="nomic-embed-text"), tools)
    graph = StateGraph(MessagesState)
    graph.add_node("agent", agent)
    graph.set_entry_point("agent")
    return graph.compile()

# ---------- 4. Interactive CLI ----------

def main():
    chroma = load_faq_to_chroma()
    agent = build_agent(chroma)
    sample_questions = [
        "What is the deadline for the assignment?",
        "How many chapters are in the course?",
        "What is the schedule for the next lecture?",
    ]
    for q in sample_questions:
        print("\nQuestion:", q)
        response = agent.invoke({"messages": [{"role": "user", "content": q}]})
        print("Answer:", response["messages"][0]["content"])
    while True:
        try:
            user_q = input("\nAsk a question (or 'exit'): ")
        except EOFError:
            break
        if user_q.lower() in {"exit", "quit"}:
            break
        response = agent.invoke({"messages": [{"role": "user", "content": user_q}]})
        print("Answer:", response["messages"][0]["content"])

if __name__ == "__main__":
    main()