"""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 os
import json
import httpx
from pathlib import Path
from dotenv import load_dotenv

from langchain_ollama import ChatOllama
from langchain_chroma import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langgraph.prebuilt import create_react_agent
from langgraph.graph import StateGraph, MessagesState

# Load environment variables (e.g., Ollama host)
load_dotenv()

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

def load_faq_to_chroma(data_dir: str = "data", persist_dir: str = "./chroma_faq"):
    """Read all .md files in data_dir, chunk them, and persist to Chroma.
    Returns the Chroma vector store.
    """
    from langchain_text_splitters import MarkdownHeaderTextSplitter
    from langchain_community.document_loaders import TextLoader
    from langchain_community.embeddings import OllamaEmbeddings

    # Ensure persist directory exists
    Path(persist_dir).mkdir(parents=True, exist_ok=True)

    # Load all markdown files
    docs = []
    for md_file in Path(data_dir).glob("*.md"):
        loader = TextLoader(str(md_file), encoding="utf-8")
        docs.extend(loader.load())

    # Split documents by headers for better context
    splitter = MarkdownHeaderTextSplitter(headers_to_split_on=["#", "##", "###"])
    split_docs = splitter.split_documents(docs)

    # Use Ollama embeddings
    embeddings = OllamaEmbeddings(model="nomic-embed-text")

    # Persist to Chroma
    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.
    """
    # For demo purposes, we load a static JSON file.
    meta_path = Path("course_meta.json")
    if not meta_path.exists():
        # Create a simple default meta file if missing
        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."""
    # Define the tool for metadata
    def meta_tool(query: str):
        return json.dumps(fetch_course_meta(query), indent=2)

    # Register tool
    tools = {"fetch_course_meta": meta_tool}

    # Prompt template with source hint
    prompt = ChatPromptTemplate.from_messages([
        ("system", "You are an FAQ bot. Use the provided tools wisely.")
    ])

    # Create a simple React agent with tool calling
    agent = create_react_agent(ChatOllama(model="llama3"), tools)

    # Graph state: messages only
    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)

    print("FAQ Bot ready. Type your question (or 'exit' to quit).")
    while True:
        user_input = input("> ")
        if user_input.lower() in {"exit", "quit"}:
            break
        # Determine if question is about metadata by simple keyword check
        if any(k in user_input.lower() for k in ["schedule", "instructor", "date"]):
            # Route to MCP tool
            response = agent.invoke({"messages": [{"role": "user", "content": user_input}]})
        else:
            # Route to Chroma via retrieval
            # Retrieve top k docs
            docs = chroma.similarity_search(user_input, k=3)
            context = "\n---\n".join(doc.page_content for doc in docs)
            # Ask the model with context
            llm = ChatOllama(model="llama3")
            answer = llm.invoke([{"role": "system", "content": "You are an FAQ bot."},
                                 {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {user_input}"}])
            print(answer.content)

if __name__ == "__main__":
    main()