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


# -----------------------------
# State
# -----------------------------
class CollegeState(TypedDict):
    query: str
    category: str
    response: str


# -----------------------------
# Intent Analyzer
# -----------------------------
def intent_analyzer(state: CollegeState):

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

    academic = [
        "admission", "apply", "jee", "eligibility",
        "exam", "semester", "course", "branch",
        "department", "syllabus"
    ]

    finance = [
        "fee", "fees", "scholarship",
        "ekalyan", "payment", "hostel fee"
    ]

    campus = [
        "hostel", "library", "canteen",
        "wifi", "transport", "placement",
        "club", "sports", "lab"
    ]

    if any(word in query for word in academic):
        category = "academic"

    elif any(word in query for word in finance):
        category = "finance"

    elif any(word in query for word in campus):
        category = "campus"

    else:
        category = "general"

    return {"category": category}


# -----------------------------
# Academic Agent
# -----------------------------
def academic_agent(state: CollegeState):

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

    if "syllabus" in q:
        response = (
            "Course-wise syllabus is available on the official college website."
        )

    elif "exam" in q:
        response = (
            "The college follows the JUT semester system. "
            "Students should regularly check the official notice section "
            "for examination schedules and results."
        )

    elif "admission" in q or "jee" in q:
        response = (
            "Admission is based on JEE Main/JCECE counseling. "
            "Candidates must satisfy the eligibility criteria announced "
            "by the admission authority."
        )

    else:
        response = (
            "The college offers undergraduate engineering programs in "
            "Civil, Mechanical, Electrical, Electronics & Communication, "
            "and Computer Science Engineering."
        )

    return {"response": response}


# -----------------------------
# Finance Agent
# -----------------------------
def finance_agent(state: CollegeState):

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

    if "scholarship" in q or "ekalyan" in q:

        response = (
            "Eligible students can apply for the E-Kalyan scholarship "
            "and other government scholarship schemes."
        )

    else:

        response = (
            "Fee details are available in the official admission "
            "prospectus on the college website."
        )

    return {"response": response}


# -----------------------------
# Campus Agent
# -----------------------------
def campus_agent(state: CollegeState):

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

    if "placement" in q:

        response = (
            "The college has a Training and Placement Cell that organizes "
            "campus drives, training programs and placement activities."
        )

    elif "hostel" in q:

        response = (
            "Hostel facilities are available for students. "
            "Contact the college administration for availability and charges."
        )

    elif "library" in q:

        response = (
            "The college library provides academic books, journals "
            "and digital learning resources."
        )

    else:

        response = (
            "The campus provides academic infrastructure, laboratories, "
            "library and student facilities."
        )

    return {"response": response}


# -----------------------------
# General Agent
# -----------------------------
def general_agent(state: CollegeState):

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

    if "location" in q or "where" in q:

        response = (
            "Dumka Engineering College is located in Dumka, Jharkhand."
        )

    elif "contact" in q:

        response = (
            "Contact details are available on the official "
            "college website."
        )

    else:

        response = (
            "Please ask about admission, fees, scholarship, "
            "placements, hostel, library or academics."
        )

    return {"response": response}


# -----------------------------
# Response Formatter
# -----------------------------
def response_formatter(state: CollegeState):

    print("\n" + "=" * 60)
    print("DEC AI ASSISTANT")
    print("=" * 60)

    print(f"\nCategory : {state['category'].title()}")

    print("\nAnswer :")
    print(state["response"])

    return state


# -----------------------------
# Router
# -----------------------------
def router(state: CollegeState):
    return state["category"]


# -----------------------------
# Build Graph
# -----------------------------
workflow = StateGraph(CollegeState)

workflow.add_node("Intent Analyzer", intent_analyzer)
workflow.add_node("Academic Agent", academic_agent)
workflow.add_node("Finance Agent", finance_agent)
workflow.add_node("Campus Agent", campus_agent)
workflow.add_node("General Agent", general_agent)
workflow.add_node("Formatter", response_formatter)

workflow.add_edge(START, "Intent Analyzer")

workflow.add_conditional_edges(
    "Intent Analyzer",
    router,
    {
        "academic": "Academic Agent",
        "finance": "Finance Agent",
        "campus": "Campus Agent",
        "general": "General Agent",
    },
)

workflow.add_edge("Academic Agent", "Formatter")
workflow.add_edge("Finance Agent", "Formatter")
workflow.add_edge("Campus Agent", "Formatter")
workflow.add_edge("General Agent", "Formatter")

workflow.add_edge("Formatter", END)

app = workflow.compile()


# -----------------------------
# Main Program
# -----------------------------
print("=" * 60)
print("🎓 Dumka Engineering College AI Assistant")
print("Type 'exit' to quit.")
print("=" * 60)

while True:

    query = input("\nYou : ")

    if query.lower() == "exit":
        print("\nThank you!")
        break

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