from typing import Dict, Any, List, Optional
from langchain_core.messages import BaseMessage, HumanMessage
from langchain_core.tools import BaseTool
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode

from .graph import BaseAgent
from .state import AgentState
from tools.calculator import CalculatorTool
from tools.web_search import WebSearchTool
from tools.rest_api_tool import RestApiTool
from tools.sr_api_tool import SrApiQueryTool


class ToolAgent(BaseAgent):
    """可使用工具完成任务的代理"""

    def __init__(self, model_section: Optional[str] = None, tools: List[BaseTool] = None):
        if tools is None:
            tools = [CalculatorTool(), WebSearchTool(), RestApiTool(), SrApiQueryTool()]

        self.tools = tools
        self.tool_node = ToolNode(tools)
        super().__init__(model_section)

    def _build_graph(self) -> StateGraph:
        """构建可使用工具的图"""
        workflow = StateGraph(AgentState)

        workflow.add_node("normalize_input", self._normalize_input)
        workflow.add_node("agent", self._agent_node)
        workflow.add_node("tools", self.tool_node)

        workflow.add_edge("normalize_input", "agent")
        workflow.add_edge("tools", "agent")

        workflow.add_conditional_edges(
            "agent",
            self._should_use_tools,
            {
                "tools": "tools",
                "end": END,
            }
        )

        workflow.set_entry_point("normalize_input")

        return workflow.compile()

    def _agent_node(self, state: AgentState) -> AgentState:
        """决定是否调用工具的代理节点"""
        model_with_tools = self.model.bind_tools(self.tools)

        if state.messages:
            try:
                response = model_with_tools.invoke(state.messages)
            except Exception as e:
                error_text = str(e)
                if "tool choice" in error_text and "auto" in error_text:
                    fallback_model = self.model.bind_tools(self.tools, tool_choice="none")
                    response = fallback_model.invoke(state.messages)
                else:
                    raise
            state.messages.append(response)

        return state

    def _should_use_tools(self, state: AgentState) -> str:
        """判断是否需要使用工具"""
        last_message = state.messages[-1]

        if hasattr(last_message, 'tool_calls') and last_message.tool_calls:
            return "tools"

        return "end"

    def run(self, user_input: str, **kwargs) -> Dict[str, Any]:
        """运行工具型代理"""
        initial_state = AgentState(
            messages=[HumanMessage(content=user_input)],
            context=kwargs
        )

        result = self.graph.invoke(initial_state)

        return {
            "messages": result.get("messages", []),
            "context": result.get("context", {}),
            "tools_used": [tool.name for tool in self.tools],
            "final_step": result.get("current_step", "unknown")
        }