import os
from typing import TypedDict, Annotated, List
import operator

from langchain_core.messages import BaseMessage, HumanMessage
from langchain_huggingface import HuggingFaceEmbeddings, ChatHuggingFace
from langchain_community.vectorstores import FAISS
from langchain.tools.retriever import create_retriever_tool
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolExecutor

# --- Configuration ---
SAVE_PATH = "/data/faiss_index"
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
# Recommended to use a powerful model for agentic tasks
LLM_REPO_ID = "mistralai/Mixtral-8x7B-Instruct-v0.1"

# --- Agent State Definition ---
class AgentState(TypedDict):
    """Defines the state of the agent graph, tracking messages."""
    messages: Annotated, operator.add]

# --- Knowledge Base and Tools Setup ---
def create_agent_system():
    """
    Initializes the entire agent system, including the knowledge base retriever,
    specialized tools, and the LangGraph-based agent executor.
    """
    print("Initializing Agent System...")
    
    # 1. Load the Knowledge Base
    if not os.path.exists(SAVE_PATH):
        raise FileNotFoundError(
            f"FAISS index not found at {SAVE_PATH}. "
            "Please run knowledge_base.py first to create it."
        )
    
    embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
    vector_store = FAISS.load_local(SAVE_PATH, embeddings, allow_dangerous_deserialization=True)
    retriever = vector_store.as_retriever(search_kwargs={'k': 3})

    # 2. Create Specialized Retriever Tools for each Agent
    # The tool descriptions are crucial as they guide the Orchestrator agent.
    academic_tool = create_retriever_tool(
        retriever,
        "academic_retriever",
        "Searches for information about academics, including curriculum, study materials, timetables, RGPV links, exam papers, and Moodle/e-Library info."
    )
    
    administrative_tool = create_retriever_tool(
        retriever,
        "administrative_retriever",
        "Searches for information about college administration, including fees, scholarships, admissions, rules, regulations, and grievance policies."
    )

    campus_services_tool = create_retriever_tool(
        retriever,
        "campus_services_retriever",
        "Searches for information about campus services like library hours, bus routes, lab availability, sports facilities, and special academies (Cisco, AWS)."
    )

    student_life_tool = create_retriever_tool(
        retriever,
        "student_life_retriever",
        "Searches for information about student life, including upcoming events, clubs, cultural festivals, and how to submit complaints or raise issues."
    )

    tools = [academic_tool, administrative_tool, campus_services_tool, student_life_tool]
    tool_executor = ToolExecutor(tools)

    # 3. Initialize the LLM
    # This requires the HUGGINGFACEHUB_API_TOKEN to be set as a secret in the Space.
    llm = ChatHuggingFace(
        repo_id=LLM_REPO_ID,
        task="text-generation",
        model_kwargs={
            "max_new_tokens": 1024,
            "temperature": 0.1,
            "repetition_penalty": 1.03,
        },
    )
    
    # Bind the tools to the LLM so it knows how to call them
    llm_with_tools = llm.bind_tools(tools)

    # 4. Define the LangGraph Nodes
    def agent_node(state):
        """The primary node that invokes the LLM to decide the next action."""
        response = llm_with_tools.invoke(state["messages"])
        return {"messages": [response]}

    def tool_node(state):
        """Executes the tool called by the agent and returns the result."""
        tool_calls = state["messages"][-1].tool_calls
        tool_messages = tool_executor.batch(tool_calls)
        return {"messages": tool_messages}

    def should_continue(state):
        """Conditional edge logic: decides whether to continue or end."""
        if state["messages"][-1].tool_calls:
            return "continue"
        return "end"

    # 5. Build the Graph
    workflow = StateGraph(AgentState)
    workflow.add_node("agent", agent_node)
    workflow.add_node("tools", tool_node)
    
    workflow.set_entry_point("agent")
    
    workflow.add_conditional_edges(
        "agent",
        should_continue,
        {"continue": "tools", "end": END}
    )
    
    workflow.add_edge("tools", "agent")
    
    # 6. Compile the graph into a runnable app
    agent_executor = workflow.compile()
    print("Agent System Initialized Successfully.")
    return agent_executor

def run_query(agent_executor, query, user_info):
    """
    Runs a query through the agent system, providing user context.
    """
    # Prepend a system message with user context for better personalization
    system_message = (
        "You are the 'GGITS Digital Campus Assistant,' a helpful, professional, and reliable AI assistant. "
        f"You are currently assisting a '{user_info['role']}' named {user_info['email']}. "
        "Use the available tools to find the most relevant and accurate information from the college's knowledge base. "
        "If you cannot find an answer, state that the information is not available in your current knowledge base."
    )
    
    messages = [
        HumanMessage(content=system_message),
        HumanMessage(content=query)
    ]
    
    # The `stream` method provides real-time output as the agent works
    final_response = None
    for chunk in agent_executor.stream({"messages": messages}):
        if "messages" in chunk:
            final_response = chunk["messages"][-1]

    return final_response.content if final_response else "I'm sorry, I couldn't process your request."