"""
极简 Agent 主图 - 用 LangGraph 原生 create_react_agent + 记忆节点
"""

from langgraph.prebuilt import create_react_agent
from langgraph.graph import StateGraph, START, END
from ..state import AgentState
from ..nodes.memory_trigger import memory_trigger_node, set_mem0_client
from backend.app.logger import info, warning
from backend.app.tools import ALL_TOOLS


def build_agent_graph(
    chat_services: dict,
    mem0_client=None,
    max_steps: int = 10
):
    """
    构建包含记忆节点的 react agent 图

    Args:
        chat_services: 模型服务字典
        mem0_client: 记忆客户端（可选）
        max_steps: 最大步数限制

    Returns:
        编译好的 graph
    """
    # 获取主模型
    primary_model = chat_services.get("primary", next(iter(chat_services.values())))

    # ========== 设置全局客户端 ==========
    if mem0_client:
        set_mem0_client(mem0_client)

    # ========== 1. 初始化节点：重置步数 ==========
    async def init_state_node(state: AgentState):
        info("[Init State] 初始化状态，重置步数")
        return {
            "current_step": 0
        }

    # ========== 2. 记忆节点（可选） ==========
    retrieve_memory_node = None
    if mem0_client:
        try:
            from ..nodes.retrieve_memory import create_retrieve_memory_node
            retrieve_memory_node = create_retrieve_memory_node(mem0_client)
        except Exception as e:
            info(f"[Graph Builder] 记忆节点初始化失败: {e}")

    # ========== 3. 创建 react agent 子图 ==========
    agent_runnable = create_react_agent(primary_model, ALL_TOOLS)

    # ========== 4. 构建主图 ==========
    graph = StateGraph(AgentState)

    graph.add_node("init_state", init_state_node)
    if retrieve_memory_node:
        graph.add_node("retrieve_memory", retrieve_memory_node)
    graph.add_node("memory_trigger", memory_trigger_node)

    # 直接把 create_react_agent 的可运行对象作为节点
    graph.add_node("agent", agent_runnable)

    # ========== 边的连接 ==========
    graph.add_edge(START, "init_state")

    if retrieve_memory_node:
        graph.add_edge("init_state", "retrieve_memory")
        graph.add_edge("retrieve_memory", "memory_trigger")
    else:
        graph.add_edge("init_state", "memory_trigger")

    graph.add_edge("memory_trigger", "agent")
    graph.add_edge("agent", END)

    info("✅ [Graph Builder] 极简 Agent 图构建完成（用 create_react_agent）")
    return graph