#pip install langgraph

from typing import TypedDict
from langgraph.graph import StateGraph, END

# -----------------------------
# Define the Graph State
# -----------------------------
class AgentState(TypedDict):
    query: str
    intent: str
    response: str


# -----------------------------
# Intent Classifier
# -----------------------------
def intent_classifier(state: AgentState):
    query = state["query"].lower()

    if "admission" in query:
        state["intent"] = "admission"
    elif "exam" in query:
        state["intent"] = "exam"
    elif "fee" in query or "fees" in query:
        state["intent"] = "fees"
    elif "scholarship" in query:
        state["intent"] = "scholarship"
    else:
        state["intent"] = "unknown"

    return state


# -----------------------------
# Admission Agent
# -----------------------------
def admission_agent(state: AgentState):
    state["response"] = (
        "Admission Agent:\n"
        "Admission is open. Submit your application, documents, and entrance exam score."
    )
    return state


# -----------------------------
# Exam Agent
# -----------------------------
def exam_agent(state: AgentState):
    state["response"] = (
        "Exam Agent:\n"
        "The semester exam starts on 05 December."
    )
    return state


# -----------------------------
# Fees Agent
# -----------------------------
def fees_agent(state: AgentState):
    state["response"] = (
        "Fees Agent:\n"
        "The semester fee is ₹42,000."
    )
    return state


# -----------------------------
# Scholarship Agent
# -----------------------------
def scholarship_agent(state: AgentState):
    state["response"] = (
        "Scholarship Agent:\n"
        "Students with more than 85% marks are eligible for scholarships."
    )
    return state


# -----------------------------
# Response Agent
# -----------------------------
def response_agent(state: AgentState):
    print("\nFinal Response:")
    print(state["response"])
    return state


# -----------------------------
# Routing Function
# -----------------------------
def router(state: AgentState):
    return state["intent"]


# -----------------------------
# Build LangGraph
# -----------------------------
workflow = StateGraph(AgentState)

workflow.add_node("Intent Classifier", intent_classifier)
workflow.add_node("Admission Agent", admission_agent)
workflow.add_node("Exam Agent", exam_agent)
workflow.add_node("Fees Agent", fees_agent)
workflow.add_node("Scholarship Agent", scholarship_agent)
workflow.add_node("Response Agent", response_agent)

workflow.set_entry_point("Intent Classifier")

workflow.add_conditional_edges(
    "Intent Classifier",
    router,
    {
        "admission": "Admission Agent",
        "exam": "Exam Agent",
        "fees": "Fees Agent",
        "scholarship": "Scholarship Agent",
    },
)

workflow.add_edge("Admission Agent", "Response Agent")
workflow.add_edge("Exam Agent", "Response Agent")
workflow.add_edge("Fees Agent", "Response Agent")
workflow.add_edge("Scholarship Agent", "Response Agent")
workflow.add_edge("Response Agent", END)

graph = workflow.compile()


# -----------------------------
# Run Example
# -----------------------------
query = input("Enter your query: ")

graph.invoke(
    {
        "query": query,
        "intent": "",
        "response": ""
    }
)