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 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("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",
                                                                                                                                                                                                                                                },
                                                                                                                                                                                                                                                )

                                                                                                                                                                                                                                                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)

                                                                                                                                                                                                                                                app = workflow.compile()

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

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