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


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


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


def classify_query(state: DECState):
    q = state["query"].lower()
    
    if any(w in q for w in ("fee", "fees", "cost", "payment", "scholarship", "financial")):
        category = "fees"
    elif any(w in q for w in ("admission", "admit", "apply", "eligibility", "cutoff", "jee")):
        category = "admission"
    elif any(w in q for w in ("placement", "recruiter", "job", "package", "salary", "median")):
        category = "placement"
    elif any(w in q for w in ("hostel", "mess", "accommodation", "dorm")):
        category = "hostel"
    elif any(w in q for w in ("library", "book", "lab", "laboratory", "facility")):
        category = "infrastructure"
    elif any(w in q for w in ("exam", "semester", "timetable", "result")):
        category = "exams"
    else:
        category = "general"

    department = "general"
    for code, name in BRANCHES.items():
        if code in q or name.lower() in q:
            department = name
            break

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


def extract_department(state: DECState):
    follow_up = ""
    if state["category"] == "general" and state["department"] == "general":
        follow_up = (
            "I can help with fees, admissions, placements, hostel, "
            "infrastructure, or exams. Try asking about B.Tech CSE fees "
            "or JEE cutoff for Mechanical."
        )
    return {"follow_up": follow_up}


def fees_agent(state: DECState):
    dep = state["department"]
    fee_map = {
        "CSE": "Rs 2,65,700 total (Rs 66,425/year)",
        "ECE": "Rs 2,65,700 total (Rs 66,425/year)",
        "EEE": "Rs 2,60,000 total (Rs 65,000/year)",
        "MECH": "Rs 2,60,000 total (Rs 65,000/year)",
        "CIVIL": "Rs 2,60,000 total (Rs 65,000/year)",
    }

    base = (
        "Dumka Engineering College fees range from Rs 66,425 to Rs 66,600 per year "
        "depending on the branch. SC/ST/OBC students can get 100% fee waiver "
        "through the E-Kalyan scholarship."
    )
    if dep in fee_map:
        base += f"\n\nFor {dep}: {fee_map[dep]}."
    return {"response": base}


def admission_agent(state: DECState):
    dep = state["department"]
    base = (
        "Admission to DEC is through JEE Main or JCECE (Jharkhand Combined Entrance). "
        "You need at least 45% in 10+2 (PCM) for General category. "
        "The application process usually runs from May to June each year."
    )

    cutoff = {
        "CSE": "around 50-60 percentile in JEE Main",
        "ECE": "around 45-55 percentile",
        "EEE": "around 40-50 percentile",
        "MECH": "around 40-50 percentile",
        "CIVIL": "around 35-45 percentile",
    }

    if dep in cutoff:
        base += f"\n\nFor {dep}, the typical cutoff is {cutoff[dep]}."
    return {"response": base}


def placement_agent(state: DECState):
    dep = state["department"]
    base = (
        "DEC has a median placement package of Rs 3.75 LPA. "
        "Top recruiters include CINIF Technologies and other IT/engineering firms. "
        "The Training & Placement cell actively prepares students with mock interviews "
        "and soft skills training."
    )

    dep_placement = {
        "CSE": "\nCSE students have the highest placement rate, with some offers up to Rs 6-8 LPA.",
        "ECE": "\nECE graduates are recruited in core electronics firms and IT companies.",
        "MECH": "\nMechanical students get opportunities in manufacturing and automotive sectors.",
        "EEE": "\nEEE students are placed in power sector companies and core electrical firms.",
        "CIVIL": "\nCivil graduates work with construction firms and government infrastructure projects.",
    }

    if dep in dep_placement:
        base += dep_placement[dep]
    return {"response": base}


def hostel_agent(state: DECState):
    return {
        "response": (
            "DEC provides separate hostels for boys and girls on campus. "
            "Hostel fee is approximately Rs 21,000 per year. "
            "Mess fee is around Rs 3,000 per month. "
            "Rooms are available on a sharing basis with basic furniture and 24/7 electricity."
        )
    }


def infrastructure_agent(state: DECState):
    return {
        "response": (
            "DEC has a well-stocked library with technical journals and reference books. "
            "The campus has dedicated labs for each department - CSE lab, ECE lab, "
            "Mechanical workshop, Civil surveying lab, and Electrical machines lab. "
            "Wi-Fi is available across campus. The college also has a sports ground "
            "and an annual cultural fest."
        )
    }


def exams_agent(state: DECState):
    return {
        "response": (
            "DEC follows the semester system under Jharkhand University of Technology (JUT), Ranchi. "
            "Exams are held twice a year - mid-semester and end-semester. "
            "Exam fee is Rs 2,000 per semester. "
            "Results are published on the college website and notice board."
        )
    }


def general_agent(state: DECState):
    return {
        "response": (
            "Dumka Engineering College, established in 2013, is located in Dumka, "
            "the second capital of Jharkhand. It offers B.Tech in 5 branches - "
            "Computer Science, Electronics & Communication, Electrical, Mechanical, "
            "and Civil Engineering. The college also offers BCA. "
            "It is affiliated with Jharkhand University of Technology and approved by AICTE."
        )
    }


def format_response(state: DECState):
    header = f"[{state['category'].upper()}]"
    if state.get("department", "general") != "general":
        header += f" - {state['department']}"

    lines = [f"\n{'='*50}", header, f"{'='*50}"]
    lines.append(state.get("response", ""))
    
    follow_up = state.get("follow_up", "")
    if follow_up:
        lines.append(f"\n-> {follow_up}")

    return {"response": "\n".join(lines)}


def route_category(state: DECState):
    return state["category"]


workflow = StateGraph(DECState)

workflow.add_node("Classifier", classify_query)
workflow.add_node("Extract Dept", extract_department)
workflow.add_node("Fees", fees_agent)
workflow.add_node("Admission", admission_agent)
workflow.add_node("Placement", placement_agent)
workflow.add_node("Hostel", hostel_agent)
workflow.add_node("Infrastructure", infrastructure_agent)
workflow.add_node("Exams", exams_agent)
workflow.add_node("General", general_agent)
workflow.add_node("Formatter", format_response)

workflow.add_edge(START, "Classifier")
workflow.add_edge("Classifier", "Extract Dept")

workflow.add_conditional_edges(
    "Extract Dept",
    route_category,
    {
        "fees": "Fees",
        "admission": "Admission",
        "placement": "Placement",
        "hostel": "Hostel",
        "infrastructure": "Infrastructure",
        "exams": "Exams",
        "general": "General",
    },
)

for agent in ("Fees", "Admission", "Placement", "Hostel", "Infrastructure", "Exams", "General"):
    workflow.add_edge(agent, "Formatter")

workflow.add_edge("Formatter", END)
app = workflow.compile()


def interactive():
    print("===== Dumka Engineering College Assistant =====")
    print("Ask about: fees, admission, placements, hostel, labs, exams, or anything about DEC.")
    print("Type 'exit' to quit.\n")

    while True:
        query = input("You: ").strip()
        if query.lower() in ("exit", "quit"):
            break
        if not query:
            continue

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

        print(result.get("response", ""))
        print()


if __name__ == "__main__":
    interactive()