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

# Reducer function jo history ko sahi se manage aur append karega
def append_history(old_list: List[str], new_list: List[str]) -> List[str]:
    return old_list + new_list

# 1. State definition fixed using LangGraph updates pattern
class UniversityState(TypedDict):
    query: str
    intent: str
    department: str
    response: str
    history: Annotated[List[str], append_history]

# Better keyword patterns matching (taaki false positives na ho)
DEPARTMENTS = {
    r"\b(cse|computer|cs)\b": "Computer Science Engineering",
    r"\b(ece|electronics)\b": "Electronics and Communication Engineering",
    r"\b(eee|electrical)\b": "Electrical Engineering",
    r"\b(me|mech|mechanical)\b": "Mechanical Engineering",
    r"\b(civil)\b": "Civil Engineering",
}

def classifier(state: UniversityState):
    text = state["query"].lower()
    
    if any(word in text for word in ["admission", "admit", "apply", "eligibility", "cutoff"]):
        intent = "admission"
    elif any(word in text for word in ["fee", "fees", "payment", "tuition"]):
        intent = "fees"
    elif any(word in text for word in ["exam", "semester", "result", "admit card", "routine"]):
        intent = "exam"
    elif any(word in text for word in ["scholarship", "e-kalyan", "financial"]):
        intent = "scholarship"
    else:
        intent = "unknown"
        
    return {"intent": intent}

def department_detector(state: UniversityState):
    text = state["query"].lower()
    department = "General"
    
    for pattern, value in DEPARTMENTS.items():
        if re.search(pattern, text):
            department = value
            break
            
    return {"department": department}

def admission_agent(state: UniversityState):
    dept = state["department"]
    response = f"""
Admission Agent
Department : {dept}

Admission Process
1. Fill Online Application Form
2. Upload Documents
   • 10th & 12th Marksheet
   • Identity Proof / ID Documents
   • Passport Size Photo
3. Pay Registration Fee
4. Submit Application
5. Download Acknowledgement
6. Wait for Merit List
7. Document Verification
8. Final Admission Confirmation
"""
    return {"response": response}

def fees_agent(state: UniversityState):
    dept = state["department"]
    response = f"""
Fees Agent
Department : {dept}

Approx Fees
Tuition Fee      : ₹73,000
Hostel Fee       : ₹20,000
Registration Fee : ₹5,000
Library Fee      : ₹2,000
Exam Fee         : ₹2,500
"""
    return {"response": response}

def exam_agent(state: UniversityState):
    response = """
Exam Agent
Semester Exam
Mid Semester : September
End Semester : December
Admit Card   : Available before exam
Result       : Usually published within 30 days.
"""
    return {"response": response}

def scholarship_agent(state: UniversityState):
    response = """
Scholarship Agent
Available Scholarships
• Merit Scholarship
• SC / ST / OBC Scholarship
• E-Kalyan Scholarship
• NSP Scholarship
"""
    return {"response": response}

def unknown_agent(state: UniversityState):
    return {
        "response": """
❌ Sorry I can't understand what you are trying to say.
I can answer only about:
• Admission
• Fees
• Exam
• Scholarship
So please ask me any of these.
"""
    }

def formatter(state: UniversityState):
    # History generator jo direct new entries update karega
    return {
        "history": [f"User : {state['query']}", f"Bot  : {state['response']}"],
        "response": state["response"]
    }

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

# --- Workflow Setup ---
workflow = StateGraph(UniversityState)

workflow.add_node("Classifier", classifier)
workflow.add_node("Department", department_detector)
workflow.add_node("Admission", admission_agent)
workflow.add_node("Fees", fees_agent)
workflow.add_node("Exam", exam_agent)
workflow.add_node("Scholarship", scholarship_agent)
workflow.add_node("Unknown", unknown_agent)
workflow.add_node("Formatter", formatter)

workflow.add_edge(START, "Classifier")
workflow.add_edge("Classifier", "Department")

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

workflow.add_edge("Admission", "Formatter")
workflow.add_edge("Fees", "Formatter")
workflow.add_edge("Exam", "Formatter")
workflow.add_edge("Scholarship", "Formatter")
workflow.add_edge("Unknown", "Formatter")
workflow.add_edge("Formatter", END)

app = workflow.compile()

def chat():
    print("=" * 60)
    print(" UNIVERSITY AI MULTI-AGENT SYSTEM")
    print("=" * 60)
    print("\nAvailable Topics: Admission, Fees, Exam, Scholarship")
    print("Departments: CSE, ECE, EEE, Mechanical, Civil")
    print("------------------------------------------------------------")
    print("Type 'bye' to exit.\n")

    # Local runtime configuration state manage karne ke liye
    chat_history = []

    while True:
        query = input("You : ").strip()
        if query.lower() in ["bye", "exit", "quit"]:
            print("\nBot : Thank you. Have a nice day!")
            break

        if not query:
            continue

        state = {
            "query": query,
            "intent": "",
            "department": "",
            "response": "",
            "history": chat_history
        }

        result = app.invoke(state)
        chat_history = result["history"]

        print("\nBot :")
        print(result["response"])
        print("-" * 40)

if __name__ == "__main__":
    chat()