import os
from typing import TypedDict, Annotated, Sequence
import operator
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode, tools_condition
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
from langchain.agents import Tool
from langchain_core.tools import tool

# Load environment variables (compatible with Hugging Face Space)
google_api_key = os.getenv("GOOGLE_API_KEY") or os.environ.get("GOOGLE_API_KEY")
if not google_api_key:
    raise ValueError("Missing GOOGLE_API_KEY environment variable")

# --- System Prompt ---
with open("System_Prompt.txt", "r", encoding="utf-8") as f:
    system_prompt = f.read()
sys_msg = SystemMessage(content=system_prompt)

# --- Tool Definitions ---
@tool
def multiply(a: int, b: int) -> int:
    """Multiply two integers together."""
    return a * b

@tool
def add(a: int, b: int) -> int:
    """Add two integers together."""
    return a + b

@tool
def subtract(a: int, b: int) -> int:
    """Subtract b from a."""
    return a - b

@tool
def divide(a: int, b: int) -> float:
    """Divide a by b. Returns float. Raises error if b is zero."""
    if b == 0:
        raise ValueError("Cannot divide by zero.")
    return a / b

@tool
def wiki_search(query: str) -> str:
    """Search Wikipedia and return up to 2 relevant documents."""
    docs = WikipediaLoader(query=query, load_max_docs=2).load()
    if not docs:
        return "No Wikipedia results found."
    return "\n\n".join([d.page_content[:1000] for d in docs])

# Tool inventory with proper categorization
tools = [
    Tool(name="Math_Multiply", func=multiply, description="Multiplies two integers"),
    Tool(name="Math_Add", func=add, description="Adds two integers"),
    Tool(name="Math_Subtract", func=subtract, description="Subtracts two integers"),
    Tool(name="Math_Divide", func=divide, description="Divides two numbers"),
    Tool(name="Search_Wikipedia", func=wiki_search, description="Searches Wikipedia"),
    Tool(
        name="Search_Web", 
        func=DuckDuckGoSearchRun().run,
        description="Searches the web using DuckDuckGo"
    )
]

# --- Graph Definition ---
class AgentState(TypedDict):
    """State definition for the agent workflow"""
    messages: Annotated[Sequence[BaseMessage], operator.add]

def build_graph():
    """Constructs and compiles the LangGraph workflow"""
    
    # Initialize LLM with Gemini 2.0 Flash
    llm = ChatGoogleGenerativeAI(
        model="gemini-2.0-flash-exp",
        temperature=0.3,
        google_api_key=google_api_key
    )
    llm_with_tools = llm.bind_tools(tools)

    # Node definitions
    def agent_node(state: AgentState):
        """Main agent node that processes messages"""
        response = llm_with_tools.invoke(state["messages"])
        return {"messages": [response]}

    # Graph construction
    workflow = StateGraph(AgentState)
    
    # Add nodes to the workflow
    workflow.add_node("agent", agent_node)
    workflow.add_node("tools", ToolNode(tools))
    
    # Configure graph flow
    workflow.set_entry_point("agent")
    workflow.add_conditional_edges(
        "agent",
        tools_condition,  # LangGraph's built-in tool detection
        {"tools": "tools", "end": END}  # Fixed END reference
    )
    workflow.add_edge("tools", "agent")
    
    return workflow.compile()

# Initialize the agent graph
agent_graph = build_graph()