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

class StudentState(TypedDict):
    query: str
    intent: str
    response: str

def intent_classifier(state: StudentState):
    question = state["query"].lower()

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

    return state

def admission_agent(state: StudentState):
    state["response"] = (
        "Admission process starts in June. Students should submit the application form with all required documents."
    )
    return state

def exam_agent(state: StudentState):
    state["response"] = (
        "The examination timetable is available on the college portal."
    )
    return state

def fees_agent(state: StudentState):
    state["response"] = (
        "Students can pay their fees online through the student portal or at the accounts office."
    )
    return state

def scholarship_agent(state: StudentState):
    state["response"] = (
        "Eligible students can apply for scholarships before the last date by submitting all required documents."
    )
    return state

def unknown_agent(state: StudentState):
    state["response"] = (
        "Sorry! I couldn't understand your question. Please ask about admission, exam, fees, or scholarship."
    )
    return state

def response_agent(state: StudentState):
    print("\n-----------------------------")
    print("Final Response")
    print("-----------------------------")
    print(state["response"])
    return state

def route(state: StudentState):
    return state["intent"]

workflow = StateGraph(StudentState)

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("Unknown Agent", unknown_agent)
workflow.add_node("Response Agent", response_agent)

workflow.set_entry_point("Intent Classifier")

workflow.add_conditional_edges(
    "Intent Classifier",
    route,
    {
        "admission": "Admission Agent",
        "exam": "Exam Agent",
        "fees": "Fees Agent",
        "scholarship": "Scholarship Agent",
        "unknown": "Unknown 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("Unknown Agent", "Response Agent")
workflow.add_edge("Response Agent", END)

app = workflow.compile()

print("===== College Student Assistant =====")

query = input("Enter your question: ")

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