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


class CollegeState(TypedDict):
    query: str
    category: str
    department: str
    response: str
    follow_up: str


BRANCHES = {
    "cse": "CSE",
    "ece": "ECE",
    "eee": "EEE",
    "mech": "MECH",
    "civil": "CIVIL",
}

def classify(state: CollegeState):

    query = state["query"].lower()

    if any(x in query for x in ["fee", "fees", "payment", "cost"]):
        category = "fees"

    elif "scholarship" in query:
        category = "scholarship"

    elif any(x in query for x in ["admission", "apply", "jee", "eligibility", "cutoff"]):
        category = "admission"

    elif any(x in query for x in ["placement", "salary", "package", "job", "recruiter"]):
        category = "placement"

    elif any(x in query for x in ["hostel", "mess", "room"]):
        category = "hostel"

    elif any(x in query for x in ["library", "lab", "wifi", "facility"]):
        category = "infrastructure"

    elif any(x in query for x in ["exam", "semester", "result", "routine"]):
        category = "exam"

    else:
        category = "general"

    department = "General"

    for short, full in BRANCHES.items():
        if short in query:
            department = full
            break

    return {
        "category": category,
        "department": department
    }


def followup(state: CollegeState):

    msg = ""

    if state["category"] == "general":
        msg = (
            "You can ask about Admission, Fees, Scholarship, "
            "Placement, Hostel, Infrastructure or Exam."
        )

    return {"follow_up": msg}

def admission_agent(state: CollegeState):

    return {
        "response":
        "Admission is through JEE Main/JCECE counselling.\n"
        "Minimum 45% marks in PCM are required."
    }


def fees_agent(state: CollegeState):

    fees = {
        "CSE": "₹66,425/year",
        "ECE": "₹66,425/year",
        "EEE": "₹65,000/year",
        "MECH": "₹65,000/year",
        "CIVIL": "₹65,000/year",
    }

    text = "Annual college fees range between ₹65,000 - ₹66,425."

    if state["department"] in fees:
        text += f"\n\n{state['department']} Fees : {fees[state['department']]}"

    return {"response": text}


def scholarship_agent(state: CollegeState):

    return {
        "response":
        "Eligible students can apply for scholarships through the "
        "state scholarship portal before the last date by submitting "
        "all required documents."
    }


def placement_agent(state: CollegeState):

    return {
        "response":
        "Median Package : ₹3.75 LPA\n"
        "Highest Package : Around ₹6-8 LPA\n"
        "Recruiters : IT and Core Companies."
    }


def hostel_agent(state: CollegeState):

    return {
        "response":
        "Hostel Fee : ₹21,000/year\n"
        "Mess Fee : ₹3,000/month"
    }


def infrastructure_agent(state: CollegeState):

    return {
        "response":
        "Facilities Available:\n"
        "- Library\n"
        "- WiFi Campus\n"
        "- Computer Labs\n"
        "- Civil Lab\n"
        "- Mechanical Workshop\n"
        "- Electrical Lab"
    }


def exam_agent(state: CollegeState):

    return {
        "response":
        "Semester System\n"
        "Exam Fee : ₹2000\n"
        "Results are published on the college website."
    }


def general_agent(state: CollegeState):

    return {
        "response":
        "Dumka Engineering College was established in 2013.\n"
        "Departments offered:\n"
        "• CSE\n"
        "• ECE\n"
        "• EEE\n"
        "• Mechanical\n"
        "• Civil"
    }

def formatter(state: CollegeState):

    title = state["category"].upper()

    if state["department"] != "General":
        title += f" ({state['department']})"

    output = []
    output.append("=" * 50)
    output.append(title)
    output.append("=" * 50)
    output.append(state["response"])

    if state["follow_up"]:
        output.append("\nSuggestion:")
        output.append(state["follow_up"])

    return {
        "response": "\n".join(output)
    }
def router(state: CollegeState):
    return state["category"]

graph = StateGraph(CollegeState)

graph.add_node("Classifier", classify)
graph.add_node("FollowUp", followup)

graph.add_node("Admission", admission_agent)
graph.add_node("Fees", fees_agent)
graph.add_node("Scholarship", scholarship_agent)
graph.add_node("Placement", placement_agent)
graph.add_node("Hostel", hostel_agent)
graph.add_node("Infrastructure", infrastructure_agent)
graph.add_node("Exam", exam_agent)
graph.add_node("General", general_agent)

graph.add_node("Formatter", formatter)

graph.add_edge(START, "Classifier")
graph.add_edge("Classifier", "FollowUp")

graph.add_conditional_edges(
    "FollowUp",
    router,
    {
        "admission": "Admission",
        "fees": "Fees",
        "scholarship": "Scholarship",
        "placement": "Placement",
        "hostel": "Hostel",
        "infrastructure": "Infrastructure",
        "exam": "Exam",
        "general": "General",
    },
)

graph.add_edge("Admission", "Formatter")
graph.add_edge("Fees", "Formatter")
graph.add_edge("Scholarship", "Formatter")
graph.add_edge("Placement", "Formatter")
graph.add_edge("Hostel", "Formatter")
graph.add_edge("Infrastructure", "Formatter")
graph.add_edge("Exam", "Formatter")
graph.add_edge("General", "Formatter")

graph.add_edge("Formatter", END)

app = graph.compile()

def main():

    print("=" * 60)
    print("      DUMKA ENGINEERING COLLEGE AI ASSISTANT")
    print("=" * 60)

    while True:

        query = input("\nYou : ")

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

        result = app.invoke(
            {
                "query": query,
                "category": "",
                "department": "",
                "response": "",
                "follow_up": "",
            }
        )

        print("\nAssistant:\n")
        print(result["response"])


if __name__ == "__main__":
    main()