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


class StudentState(TypedDict):
    query: str
        cleaned_query: str
            intent: str
                confidence: float
                    response: str
                        history: List[str]


                        ADMISSION_KEYWORDS = [
                            "admission","apply","application","join","enroll",
                                "enrollment","eligibility","document","documents","registration"
                                ]

                                EXAM_KEYWORDS = [
                                    "exam","examination","test","result",
                                        "schedule","routine","timetable","marks"
                                        ]

                                        FEE_KEYWORDS = [
                                            "fee","fees","payment","tuition",
                                                "installment","amount","online payment"
                                                ]

                                                SCHOLARSHIP_KEYWORDS = [
                                                    "scholarship","financial aid","grant",
                                                        "stipend","funding"
                                                        ]


                                                        def preprocess_query(query: str) -> str:
                                                            query = query.lower()
                                                                query = re.sub(r"[^a-zA-Z0-9 ]", "", query)
                                                                    query = re.sub(r"\s+", " ", query)
                                                                        return query.strip()


                                                                        def calculate_confidence(keyword_count: int) -> float:
                                                                            if keyword_count >= 3:
                                                                                    return 0.98
                                                                                        elif keyword_count == 2:
                                                                                                return 0.90
                                                                                                    elif keyword_count == 1:
                                                                                                            return 0.80
                                                                                                                return 0.30


                                                                                                                def preprocessing_agent(state: StudentState):
                                                                                                                    state["cleaned_query"] = preprocess_query(state["query"])
                                                                                                                        return state


                                                                                                                        def intent_classifier(state: StudentState):
                                                                                                                            query = state["cleaned_query"]

                                                                                                                                scores = {
                                                                                                                                        "admission": sum(k in query for k in ADMISSION_KEYWORDS),
                                                                                                                                                "exam": sum(k in query for k in EXAM_KEYWORDS),
                                                                                                                                                        "fees": sum(k in query for k in FEE_KEYWORDS),
                                                                                                                                                                "scholarship": sum(k in query for k in SCHOLARSHIP_KEYWORDS),
                                                                                                                                                                    }

                                                                                                                                                                        intent = max(scores, key=scores.get)
                                                                                                                                                                            highest = scores[intent]

                                                                                                                                                                                if highest == 0:
                                                                                                                                                                                        state["intent"] = "unknown"
                                                                                                                                                                                                state["confidence"] = 0.30
                                                                                                                                                                                                    else:
                                                                                                                                                                                                            state["intent"] = intent
                                                                                                                                                                                                                    state["confidence"] = calculate_confidence(highest)

                                                                                                                                                                                                                        return state


                                                                                                                                                                                                                        def admission_agent(state):
                                                                                                                                                                                                                            state["response"] = """Admission Information

                                                                                                                                                                                                                            • Admission starts in June every year.
                                                                                                                                                                                                                            • Fill out the application form.
                                                                                                                                                                                                                            • Upload required documents.
                                                                                                                                                                                                                            • Pay the application fee.
                                                                                                                                                                                                                            • Wait for the merit list.
                                                                                                                                                                                                                            • Complete document verification."""
                                                                                                                                                                                                                                return state


                                                                                                                                                                                                                                def exam_agent(state):
                                                                                                                                                                                                                                    state["response"] = """Examination Information

                                                                                                                                                                                                                                    • Check the timetable on the student portal.
                                                                                                                                                                                                                                    • Download your admit card.
                                                                                                                                                                                                                                    • Carry your ID card.
                                                                                                                                                                                                                                    • Results are published after evaluation."""
                                                                                                                                                                                                                                        return state


                                                                                                                                                                                                                                        def fees_agent(state):
                                                                                                                                                                                                                                            state["response"] = """Fees Information

                                                                                                                                                                                                                                            • Pay fees online or at the Accounts Office.
                                                                                                                                                                                                                                            • UPI, Net Banking and Cards are supported.
                                                                                                                                                                                                                                            • Download the payment receipt.
                                                                                                                                                                                                                                            • And fees depend on the mode of admission. (paid or free shet)"""
                                                                                                                                                                                                                                                return state


                                                                                                                                                                                                                                                def scholarship_agent(state):
                                                                                                                                                                                                                                                    state["response"] = """Scholarship Information

                                                                                                                                                                                                                                                    • Eligible students can apply online.
                                                                                                                                                                                                                                                    • Submit all required documents.
                                                                                                                                                                                                                                                    • Apply before the deadline."""
                                                                                                                                                                                                                                                        return state


                                                                                                                                                                                                                                                        def unknown_agent(state):
                                                                                                                                                                                                                                                            state["response"] = """Sorry, I couldn't understand your question.

                                                                                                                                                                                                                                                            Please ask about:
                                                                                                                                                                                                                                                            • Admission
                                                                                                                                                                                                                                                            • Examination
                                                                                                                                                                                                                                                            • Fees
                                                                                                                                                                                                                                                            • Scholarship"""
                                                                                                                                                                                                                                                                return state


                                                                                                                                                                                                                                                                def router(state):
                                                                                                                                                                                                                                                                    return state["intent"]


                                                                                                                                                                                                                                                                    def save_history(state):
                                                                                                                                                                                                                                                                        state["history"].append(state["query"])
                                                                                                                                                                                                                                                                            return state


                                                                                                                                                                                                                                                                            def response_formatter(state):
                                                                                                                                                                                                                                                                                state["response"] = f"""
                                                                                                                                                                                                                                                                                ========================================
                                                                                                                                                                                                                                                                                College Student Assistant
                                                                                                                                                                                                                                                                                ========================================

                                                                                                                                                                                                                                                                                {state["response"]}

                                                                                                                                                                                                                                                                                ========================================
                                                                                                                                                                                                                                                                                """
                                                                                                                                                                                                                                                                                    return state


                                                                                                                                                                                                                                                                                    def response_agent(state):
                                                                                                                                                                                                                                                                                        print(state["response"])
                                                                                                                                                                                                                                                                                            return state


                                                                                                                                                                                                                                                                                            workflow = StateGraph(StudentState)

                                                                                                                                                                                                                                                                                            workflow.add_node("Preprocessing", preprocessing_agent)
                                                                                                                                                                                                                                                                                            workflow.add_node("Intent", intent_classifier)
                                                                                                                                                                                                                                                                                            workflow.add_node("Admission", admission_agent)
                                                                                                                                                                                                                                                                                            workflow.add_node("Exam", exam_agent)
                                                                                                                                                                                                                                                                                            workflow.add_node("Fees", fees_agent)
                                                                                                                                                                                                                                                                                            workflow.add_node("Scholarship", scholarship_agent)
                                                                                                                                                                                                                                                                                            workflow.add_node("Unknown", unknown_agent)
                                                                                                                                                                                                                                                                                            workflow.add_node("History", save_history)
                                                                                                                                                                                                                                                                                            workflow.add_node("Formatter", response_formatter)
                                                                                                                                                                                                                                                                                            workflow.add_node("Response", response_agent)

                                                                                                                                                                                                                                                                                            workflow.set_entry_point("Preprocessing")
                                                                                                                                                                                                                                                                                            workflow.add_edge("Preprocessing", "Intent")

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

                                                                                                                                                                                                                                                                                                                                                    for node in ["Admission", "Exam", "Fees", "Scholarship", "Unknown"]:
                                                                                                                                                                                                                                                                                                                                                        workflow.add_edge(node, "History")

                                                                                                                                                                                                                                                                                                                                                        workflow.add_edge("History", "Formatter")
                                                                                                                                                                                                                                                                                                                                                        workflow.add_edge("Formatter", "Response")
                                                                                                                                                                                                                                                                                                                                                        workflow.add_edge("Response", END)

                                                                                                                                                                                                                                                                                                                                                        app = workflow.compile()


                                                                                                                                                                                                                                                                                                                                                        def main():
                                                                                                                                                                                                                                                                                                                                                            print("=" * 50)
                                                                                                                                                                                                                                                                                                                                                                print("COLLEGE STUDENT ASSISTANT USING LANGGRAPH")
                                                                                                                                                                                                                                                                                                                                                                    print("=" * 50)
                                                                                                                                                                                                                                                                                                                                                                        print("Type 'exit' to quit.\n")

                                                                                                                                                                                                                                                                                                                                                                            history = []

                                                                                                                                                                                                                                                                                                                                                                                while True:
                                                                                                                                                                                                                                                                                                                                                                                        query = input("Ask your question: ").strip()

                                                                                                                                                                                                                                                                                                                                                                                                if query.lower() in ("exit", "quit"):
                                                                                                                                                                                                                                                                                                                                                                                                            print("Goodbye!")
                                                                                                                                                                                                                                                                                                                                                                                                                        break

                                                                                                                                                                                                                                                                                                                                                                                                                                if not query:
                                                                                                                                                                                                                                                                                                                                                                                                                                            print("Please enter a question.\n")
                                                                                                                                                                                                                                                                                                                                                                                                                                                        continue

                                                                                                                                                                                                                                                                                                                                                                                                                                                                state = {
                                                                                                                                                                                                                                                                                                                                                                                                                                                                            "query": query,
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                        "cleaned_query": "",
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                    "intent": "",
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                "confidence": 0.0,
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                            "response": "",
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                        "history": history,
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                }

                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                        app.invoke(state)


                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                        if __name__ == "__main__":
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                            main()