import os
from dotenv import load_dotenv

from typing import TypedDict

from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage

load_dotenv()

llm = ChatOpenAI(model="gpt-4o-mini")

# -----------------------------
# State
# -----------------------------

class AgentState(TypedDict):
    question: str
    intent: str
    answer: str


# -----------------------------
# Intent Classifier
# -----------------------------

def intent_classifier(state: AgentState):

    question = state["question"]

    prompt = f"""
You are an Intent Classifier.

Classify the question into ONLY one category.

Categories:

Admission
Exam
Fees
Scholarship

Question:
{question}

Return only category name.
"""

    response = llm.invoke([HumanMessage(content=prompt)])

    intent = response.content.strip()

    return {"intent": intent}


# -----------------------------
# Admission Agent
# -----------------------------

def admission_agent(state: AgentState):

    question = state["question"]

    prompt = f"""
You are Admission Agent.

Answer admission related questions.

Question:
{question}
"""

    response = llm.invoke(question)

    return {"answer": response.content}


# -----------------------------
# Exam Agent
# -----------------------------

def exam_agent(state: AgentState):

    question = state["question"]

    prompt = f"""
You are Exam Agent.

Answer exam related questions.

Question:
{question}
"""

    response = llm.invoke(prompt)

    return {"answer": response.content}


# -----------------------------
# Fees Agent
# -----------------------------

def fees_agent(state: AgentState):

    question = state["question"]

    prompt = f"""
You are Fees Agent.

Answer fee related questions.

Question:
{question}
"""

    response = llm.invoke(prompt)

    return {"answer": response.content}


# -----------------------------
# Scholarship Agent
# -----------------------------

def scholarship_agent(state: AgentState):

    question = state["question"]

    prompt = f"""
You are Scholarship Agent.

Answer scholarship related questions.

Question:
{question}
"""

    response = llm.invoke(prompt)

    return {"answer": response.content}


# -----------------------------
# Response Agent
# -----------------------------

def response_agent(state: AgentState):

    return {"answer": state["answer"]}


# -----------------------------
# Router
# -----------------------------

def router(state: AgentState):

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

    if "admission" in intent:
        return "Admission Agent"

    elif "exam" in intent:
        return "Exam Agent"

    elif "fees" in intent:
        return "Fees Agent"

    else:
        return "Scholarship Agent"


# -----------------------------
# Graph
# -----------------------------

graph = StateGraph(AgentState)

graph.add_node("Intent Classifier", intent_classifier)

graph.add_node("Admission Agent", admission_agent)
graph.add_node("Exam Agent", exam_agent)
graph.add_node("Fees Agent", fees_agent)
graph.add_node("Scholarship Agent", scholarship_agent)

graph.add_node("Response Agent", response_agent)

graph.set_entry_point("Intent Classifier")

graph.add_conditional_edges(
    "Intent Classifier",
    router,
    {
        "Admission Agent": "Admission Agent",
        "Exam Agent": "Exam Agent",
        "Fees Agent": "Fees Agent",
        "Scholarship Agent": "Scholarship Agent",
    },
)

graph.add_edge("Admission Agent", "Response Agent")
graph.add_edge("Exam Agent", "Response Agent")
graph.add_edge("Fees Agent", "Response Agent")
graph.add_edge("Scholarship Agent", "Response Agent")

graph.add_edge("Response Agent", END)

app = graph.compile()


# -----------------------------
# Run
# -----------------------------

while True:

    q = input("\nAsk Question: ")

    if q.lower() == "exit":
        break

    result = app.invoke(
        {
            "question": q,
            "intent": "",
            "answer": "",
        }
    )

    print("\nAnswer:\n")
    print(result["answer"])