from typing import Dict, Any, List, Optional
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, ToolMessage
from langchain_core.tools import BaseTool
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode
from .base_agent import BaseAgent, AgentState
from tools.calculator import CalculatorTool
from tools.web_search import WebSearchTool


class ToolAgent(BaseAgent):
    """Agent that can use tools to accomplish tasks"""
    
    def __init__(self, model_name: str = None, tools: List[BaseTool] = None):
        # Initialize with default tools if none provided
        if tools is None:
            tools = [CalculatorTool(), WebSearchTool()]
        
        self.tools = tools
        self.tool_node = ToolNode(tools)
        super().__init__(model_name)
    
    def _build_graph(self) -> StateGraph:
        """Build tool-using graph"""
        workflow = StateGraph(AgentState)
        
        # Add nodes
        workflow.add_node("agent", self._agent_node)
        workflow.add_node("tools", self.tool_node)
        
        # Define edges
        workflow.add_edge("tools", "agent")
        
        # Conditional routing
        workflow.add_conditional_edges(
            "agent",
            self._should_use_tools,
            {
                "tools": "tools",
                "end": END,
            }
        )
        
        # Set entry point
        workflow.set_entry_point("agent")
        
        return workflow.compile()
    
    def _agent_node(self, state: AgentState) -> AgentState:
        """Agent node that decides whether to use tools"""
        # Bind tools to the model
        model_with_tools = self.model.bind_tools(self.tools)
        
        # Get the last message
        if state.messages:
            response = model_with_tools.invoke(state.messages)
            state.messages.append(response)
        
        return state
    
    def _should_use_tools(self, state: AgentState) -> str:
        """Determine if tools should be used"""
        last_message = state.messages[-1]
        
        # If the last message has tool calls, route to tools
        if hasattr(last_message, 'tool_calls') and last_message.tool_calls:
            return "tools"
        
        # Otherwise, end the workflow
        return "end"
    
    def run(self, user_input: str, **kwargs) -> Dict[str, Any]:
        """Run the tool-using agent"""
        initial_state = AgentState(
            messages=[HumanMessage(content=user_input)],
            context=kwargs
        )
        
        result = self.graph.invoke(initial_state)
        
        return {
            "messages": result.messages,
            "context": result.context,
            "tools_used": [tool.name for tool in self.tools],
            "final_step": result.current_step
        }