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

class State(TypedDict):
    query:str
    intent:str
    answer:str
    

builder = StateGraph(State)

def intent_classifier(state):
    query=state["query"].lower()

    if "admission" in query:
        return{"intent":"admission"}

    if "exam" in query:
        return{"intent":"exam"}
    
    if "fee" in query:
        return{"intent":"fee"}

    if "scholarship" in query:
        return{"intent":"scholarship"}

    return{"intent":"unknow"}


def admission_agent(state):
    return{"answer":"Admission details......"}

def exam_agent(state):
    return{"answer":"exam details......"}

def fee_agent(state):
    return{"answer":"fee details......"}

def scholarship_agent(state):
    return{"answer":"scholarship details......"}

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

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


builder.add_node("intent" ,intent_classifier)
builder.add_node("admission" ,admission_agent)
builder.add_node("exam" ,exam_agent)
builder.add_node("fee" ,fee_agent)
builder.add_node("scholarship" ,scholarship_agent)
builder.add_node("response" ,response_agent) 


builder.add_edge(START ,"intent")

builder.add_conditional_edges(
    "intent",
    router,
    {
        "admission" : "admission",
        "exam" : "exam",
        "fee" : "fee",
        "scholarship" : "scholarship",
    }
)

builder.add_edge("admission","response")
builder.add_edge("exam","response")
builder.add_edge("fee","response")
builder.add_edge("scholarship","response")
builder.add_edge("response",END)

graph = builder.compile()


intial_state = {"query":"tell me fee off class 5"}

result = graph.invoke(intial_state)

print(result)