"""LangGraph agent that routes queries to either Chroma or the MCP tool.

The agent uses a simple decision rule: if the query contains the word
"schedule" or "metadata", it calls the MCP tool; otherwise it searches
Chroma.  The answer is annotated with a `source` field.
"""

from typing import Dict

from langgraph.graph import StateGraph, END
from langchain_core.runnables import Runnable

from .rag_tools import search_course_docs
from .tools import fetch_course_meta

# ---------- State definition ----------
class AgentState(dict):
    """State passed between nodes.

    Keys:
    - ``input``: the user question
    - ``answer``: final answer string
    - ``source``: "chroma" or "mcp_meta"
    """

    def __init__(self, input: str = ""):
        super().__init__(input=input, answer="", source="")

# ---------- Decision node ----------
async def decide_route(state: AgentState) -> AgentState:
    """Return the state unchanged – routing is handled by a conditional edge.

    The decision logic is intentionally simple to keep the example
    readable.
    """
    return state

# ---------- Chroma search node ----------
async def search_chroma(state: AgentState) -> AgentState:
    query = state["input"]
    docs = search_course_docs(query)
    # Concatenate the snippets – in a real system we might use a
    # retrieval‑augmented generation chain.
    snippet = "\n\n".join(doc.page_content for doc in docs)
    state["answer"] = snippet if snippet else "No relevant information found."
    state["source"] = "chroma"
    return state

# ---------- MCP tool node ----------
async def fetch_meta(state: AgentState) -> AgentState:
    query = state["input"]
    meta = fetch_course_meta(query)
    state["answer"] = meta if isinstance(meta, str) else str(meta)
    state["source"] = "mcp_meta"
    return state

# ---------- Final node ----------
async def finalize(state: AgentState) -> AgentState:
    # The agent already has the answer; we just return the state.
    return state

# ---------- Build the graph ----------
async def create_agent_executor() -> Runnable:
    # Build the graph
    graph = StateGraph(AgentState)
    graph.add_node("decide", decide_route)
    graph.add_node("search_chroma", search_chroma)
    graph.add_node("fetch_meta", fetch_meta)
    graph.add_node("finalize", finalize)

    graph.set_entry_point("decide")
    # Conditional edge based on the query content
    def decide_condition(state: AgentState) -> str:
        query = state["input"].lower()
        if "schedule" in query or "metadata" in query:
            return "fetch_meta"
        return "search_chroma"

    graph.add_conditional_edges("decide", decide_condition)
    graph.add_edge("search_chroma", "finalize")
    graph.add_edge("fetch_meta", "finalize")
    graph.add_edge("finalize", END)

    return graph.compile()