import os
from dotenv import load_dotenv

from typing_extensions import TypedDict

from langgraph.graph import StateGraph, START, END

from langchain_google_genai import ChatGoogleGenerativeAI


load_dotenv()

GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")

if not GOOGLE_API_KEY:
    raise ValueError(
        "GOOGLE_API_KEY not found.\n"
        "Create a .env file and add:\n"
        "GOOGLE_API_KEY=YOUR_API_KEY"
    )



class GeminiLLM:
    """
    Wrapper class for Google Gemini.
    """

    def __init__(self):

        self.model = ChatGoogleGenerativeAI(
            model="gemini-3.5-flash",
            google_api_key=GOOGLE_API_KEY,
            temperature=0.3,
        )

    def generate(self, prompt: str) -> str:

        try:

            response = self.model.invoke(prompt)

            if isinstance(response.content, str):
                return response.content

            if isinstance(response.content, list):

                text = ""

                for block in response.content:

                    if (
                        isinstance(block, dict)
                        and block.get("type") == "text"
                    ):
                        text += block.get("text", "")

                return text.strip()

            return str(response.content)

        except Exception as e:

            return f"Error: {e}"


gemini = GeminiLLM()



# STATE


class CollegeState(TypedDict):

    query: str

    intent: str

    response: str



# PYTHON INTENT CLASSIFIER
# (NO GEMINI API CALL)


def classify_intent(query: str) -> str:

    q = query.lower()

    admission_keywords = [
        "admission",
        "apply",
        "application",
        "eligibility",
        "document",
        "admit",
        "counselling",
        "cutoff",
    ]

    exam_keywords = [
        "exam",
        "result",
        "attendance",
        "semester",
        "internal",
        "backlog",
        "syllabus",
        "schedule",
    ]

    fees_keywords = [
        "fee",
        "fees",
        "hostel fee",
        "payment",
        "refund",
        "tuition",
        "bus fee",
    ]

    scholarship_keywords = [
        "scholarship",
        "nsp",
        "merit",
        "government scholarship",
        "state scholarship",
    ]

    for word in admission_keywords:

        if word in q:
            return "admission"

    for word in exam_keywords:

        if word in q:
            return "exam"

    for word in fees_keywords:

        if word in q:
            return "fees"

    for word in scholarship_keywords:

        if word in q:
            return "scholarship"

    return "unknown"



# INTENT NODE


def intent_classifier(state: CollegeState):

    state["intent"] = classify_intent(state["query"])

    return state



# ROUTER


def router(state: CollegeState):

    return state["intent"]



# UNKNOWN AGENT


def unknown_agent(state: CollegeState):

    state["response"] = (
        "Sorry, I can only answer questions related to:\n\n"
        "• Admission\n"
        "• Examination\n"
        "• Fees\n"
        "• Scholarship"
    )

    return state

def exam_agent(state: CollegeState):

    prompt = f"""
You are a college exam assistant.

Answer only examination-related questions.

Question:
{state["query"]}
"""

    state["response"] = gemini.generate(prompt)

    return state


# ADMISSION AGENT

def admission_agent(state: CollegeState):

    prompt = f"""
You are a college admission assistant.

Answer only admission-related questions in simple language.

Question:
{state["query"]}
"""

    state["response"] = gemini.generate(prompt)

    return state



# FEES AGENT


def fees_agent(state: CollegeState):
    prompt = f"""
You are a college fees assistant.

Answer only fee-related questions.

Question:
{state["query"]}
"""

    state["response"] = gemini.generate(prompt)

    return state



# SCHOLARSHIP AGENT


def scholarship_agent(state: CollegeState):

    prompt = f"""
You are a scholarship assistant.

Answer only scholarship-related questions.

Question:
{state["query"]}
"""

    state["response"] = gemini.generate(prompt)

    return state




def response_agent(state: CollegeState):
    """
    Optional response formatter.

    No additional Gemini API call.
    This keeps the project faster and avoids
    consuming extra API quota.
    """

    response = state["response"].strip()

    if not response:
        response = "No response generated."

    state["response"] = response

    return state

# BUILD LANGGRAPH


builder = StateGraph(CollegeState)


# Register Nodes


builder.add_node("intent_classifier", intent_classifier)
builder.add_node("admission", admission_agent)
builder.add_node("exam", exam_agent)
builder.add_node("fees", fees_agent)
builder.add_node("scholarship", scholarship_agent)
builder.add_node("unknown", unknown_agent)
builder.add_node("response", response_agent)


# START


builder.add_edge(START, "intent_classifier")


# Conditional Routing

builder.add_conditional_edges(
    "intent_classifier",
    router,
    {
        "admission": "admission",
        "exam": "exam",
        "fees": "fees",
        "scholarship": "scholarship",
        "unknown": "unknown",
    },
)


builder.add_edge("admission", "response")
builder.add_edge("exam", "response")
builder.add_edge("fees", "response")
builder.add_edge("scholarship", "response")
builder.add_edge("unknown", "response")

# -----------------------------
# Response -> END
# -----------------------------

builder.add_edge("response", END)

# Compile Graph
graph = builder.compile()



def main():

    print("=" * 60)
    print(" College Information Multi-Agent AI agent ")
    print("=" * 60)
    print("Type 'exit' to quit.\n")

    while True:

        query = input(" Student: ").strip()

        if query.lower() in ["exit", "quit"]:

            print("\n Thank you budy")
            break

        state = {
            "query": query,
            "intent": "",
            "response": "",
        }

        try:

            result = graph.invoke(state)

            print("\n----------------------------------------")
            print(f"Detected Intent : {result['intent']}")
            print("----------------------------------------")

            print("\n Agent:\n")
            print(result["response"])
            print()

        except Exception as e:

            print("\n❌ Error")
            print(e)
            print()


if __name__ == "__main__":
    main()