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


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


def classifier(state: DEC):
    q = state["query"].lower()

    if any(w in q for w in ("fee", "fees", "cost")):
        intent = "fees"
    elif any(w in q for w in ("exam", "semester", "timetable", "result")):
        intent = "exam"
    elif any(w in q for w in ("admission", "admit", "apply", "eligibility", "cutoff", "jee")):
        intent = "admission"
    elif any(w in q for w in ("scholarship", "ekalyan", "financial")):
        intent = "scholarship"
    else:
        intent = "unknown"

    return {"intent": intent}


def fees_agent(state: DEC):
    return {
        "response": (
            "DEC fees are Rs 66,425-66,600 per year depending on branch. "
            "Total B.Tech fee ranges from Rs 2.6L to Rs 2.65L. "
            "SC/ST/OBC students can get 100% waiver via E-Kalyan."
        )
    }


def exam_agent(state: DEC):
    return {
        "response": (
            "DEC follows the JUT semester system. Exams are held twice a year "
            "(mid-sem and end-sem). Exam fee is Rs 2,000 per semester. "
            "Results are posted on the college website and notice board."
        )
    }


def admission_agent(state: DEC):
    return {
        "response": (
            "Admission is through JEE Main or JCECE. Min 45% in 10+2 (PCM) required. "
            "Application runs May-June. CSE cutoff is ~50-60 percentile, "
            "ECE ~45-55, EEE/MECH ~40-50, Civil ~35-45."
        )
    }


def scholarship_agent(state: DEC):
    return {
        "response": (
            "DEC students can apply for E-Kalyan scholarship (100% fee waiver for "
            "SC/ST/OBC). Other merit-based scholarships are available. "
            "Apply before the deadline with all required documents."
        )
    }


def unknown_agent(state: DEC):
    return {
        "response": "I couldn't understand. Try asking about fees, exams, admission, or scholarship."
    }


def response_agent(state: DEC):
    print("\n" + "="*50)
    print(f"[{state['intent'].upper()}]")
    print("="*50)
    print(state.get("response", ""))
    return state


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


workflow = StateGraph(DEC)

workflow.add_node("Classifier", classifier)
workflow.add_node("Fees", fees_agent)
workflow.add_node("Exam", exam_agent)
workflow.add_node("Admission", admission_agent)
workflow.add_node("Scholarship", scholarship_agent)
workflow.add_node("Unknown", unknown_agent)
workflow.add_node("Response", response_agent)

workflow.add_edge(START, "Classifier")

workflow.add_conditional_edges(
    "Classifier",
    route,
    {
        "fees": "Fees",
        "exam": "Exam",
        "admission": "Admission",
        "scholarship": "Scholarship",
        "unknown": "Unknown",
    },
)

for agent in ("Fees", "Exam", "Admission", "Scholarship", "Unknown"):
    workflow.add_edge(agent, "Response")

workflow.add_edge("Response", END)

app = workflow.compile()


print("===== Dumka Engineering College Assistant =====")
query = input("Enter your question: ")

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