"""
OpenCode Style agent 的图路由与条件边构建。

本模块从 `factory.py` 中拆分出与具体业务无关的 LangGraph 结构逻辑：
- 基于 OpenCode 状态字段（例如 `should_end`、`structured_response`）决定下一跳节点
- 构建 model ↔ tools ↔ middleware 之间的条件边
- 封装 middleware `jump_to` 语义到统一的 `_resolve_jump` 辅助函数

与 LangChain 官方 agent 的差异在于：
- 这里的边逻辑假设退出由 OpenCode Style 的状态/finish_reason 控制，
  而不是传统的“是否还有 tool_calls 或 recursion_limit”模式。
"""

from __future__ import annotations

from typing import Any, Callable, List, Tuple

from langchain.agents.middleware.types import (
    AgentState,
    ContextT,
    JumpTo,
    ResponseT,
    _InputAgentState,
    _OutputAgentState,
)
from langchain.agents.structured_output import OutputToolBinding
from langchain_core.messages import AIMessage, AnyMessage, ToolMessage
from langgraph._internal._runnable import RunnableCallable
from langgraph.graph.state import StateGraph
from langgraph.prebuilt.tool_node import ToolCallWithContext, ToolNode
from langgraph.types import Send

from .exit_logic import MAX_STEPS_PROMPT


def _resolve_jump(
    jump_to: JumpTo | None,
    *,
    model_destination: str,
    end_destination: str,
) -> str | None:
    if jump_to == "model":
        return model_destination
    if jump_to == "end":
        return end_destination
    if jump_to == "tools":
        return "tools"
    return None


def _fetch_last_ai_and_tool_messages(
    messages: List[AnyMessage],
) -> Tuple[AIMessage | None, List[ToolMessage]]:
    """Return the last AI message and any subsequent tool messages."""
    for i in range(len(messages) - 1, -1, -1):
        if isinstance(messages[i], AIMessage):
            last_ai_message = messages[i]
            tool_messages = [m for m in messages[i + 1 :] if isinstance(m, ToolMessage)]
            return last_ai_message, tool_messages

    return None, []


def _make_loop_controller_edge(
    *,
    model_destination: str,
    tools_destination: str,
    structured_output_tools: dict[str, OutputToolBinding[Any]],
    end_destination: str,
) -> Callable[[dict[str, Any]], str | List[Send] | None]:
    """
    OpenCode Style 的 loop_controller 条件边。

    职责：
    - 优先处理 middleware 注入的 `jump_to`
    - 处理 should_end / end_reason（包括 max_steps_reached 特殊分支）
    - 决定是继续走 tools、回到 model，还是直接结束
    """

    def loop_controller(
        state: dict[str, Any],
    ) -> str | List[Send] | None:
        # 1. 如果 middleware 显式指定了 jump_to，则优先遵循
        if jump_to := state.get("jump_to"):
            return _resolve_jump(
                jump_to,
                model_destination=model_destination,
                end_destination=end_destination,
            )

        should_end = bool(state.get("should_end"))
        end_reason = state.get("end_reason")

        last_ai_message, tool_messages = _fetch_last_ai_and_tool_messages(state["messages"])

        # 2. 如果 model 判定应该结束
        if should_end:
            # 2.a 最大步数打满：追加 MAX_STEPS_PROMPT，清理结束标记，再回到 model 做最后一次“纯文本总结”
            if end_reason == "max_steps_reached":
                messages = list(state.get("messages") or [])
                messages.append(AIMessage(content=MAX_STEPS_PROMPT))
                state["messages"] = messages

                # 清理结束标记，让下一轮 model 正常运行
                state["should_end"] = False
                state["end_reason"] = None

                # 可以选择加一个标记，方便后续调试/监控
                state["force_text_exit"] = True

                return model_destination

            # 2.b 其他结束原因：直接结束
            return end_destination

        # 3. should_end=False：按常规 OpenCode Style 逻辑判断是否还有工具调用
        if last_ai_message is None:
            # 如果连 AIMessage 都没有，说明状态可能被清空，直接结束
            return end_destination

        tool_message_ids = [m.tool_call_id for m in tool_messages]

        pending_tool_calls = [
            c
            for c in last_ai_message.tool_calls
            if c["id"] not in tool_message_ids and c["name"] not in structured_output_tools
        ]

        if pending_tool_calls:
            return [
                Send(
                    tools_destination,
                    ToolCallWithContext(
                        __type="tool_call_with_context",
                        tool_call=tool_call,
                        state=state,
                    ),
                )
                for tool_call in pending_tool_calls
            ]

        # 4. 如果有结构化输出，直接结束
        if "structured_response" in state:
            return end_destination

        # 5. 没有待处理的工具调用，结束
        return end_destination

    return loop_controller


def _make_model_to_tools_edge(
    *,
    model_destination: str,
    structured_output_tools: dict[str, OutputToolBinding[Any]],
    end_destination: str,
) -> Callable[[dict[str, Any]], str | List[Send] | None]:
    """OpenCode Style 的 model → tools/END 条件边（基于 should_end 而非 tool_calls 判断）。"""

    def model_to_tools(
        state: dict[str, Any],
    ) -> str | List[Send] | None:
        # 1. If there's an explicit jump_to in the state, use it
        if jump_to := state.get("jump_to"):
            return _resolve_jump(
                jump_to,
                model_destination=model_destination,
                end_destination=end_destination,
            )

        # OpenCode Style: 基于 should_end 判断
        if state.get("should_end"):
            return end_destination

        last_ai_message, tool_messages = _fetch_last_ai_and_tool_messages(state["messages"])

        # 2. if no AIMessage exists (e.g., messages were cleared), exit the loop
        if last_ai_message is None:
            return end_destination

        tool_message_ids = [m.tool_call_id for m in tool_messages]

        # 3. 如果有未执行的工具调用，跳转到工具节点
        pending_tool_calls = [
            c
            for c in last_ai_message.tool_calls
            if c["id"] not in tool_message_ids and c["name"] not in structured_output_tools
        ]

        if pending_tool_calls:
            return [
                Send(
                    "tools",
                    ToolCallWithContext(
                        __type="tool_call_with_context",
                        tool_call=tool_call,
                        state=state,
                    ),
                )
                for tool_call in pending_tool_calls
            ]

        # 4. If there is a structured response, exit the loop
        if "structured_response" in state:
            return end_destination

        # 5. 没有待处理的工具调用，结束
        return end_destination

    return model_to_tools


def _make_model_to_model_edge(
    *,
    model_destination: str,
    end_destination: str,
) -> Callable[[dict[str, Any]], str | List[Send] | None]:
    """model → model/END 条件边，用于仅依赖结构化输出的场景。"""

    def model_to_model(
        state: dict[str, Any],
    ) -> str | List[Send] | None:
        # 1. Priority: Check for explicit jump_to directive from middleware
        if jump_to := state.get("jump_to"):
            return _resolve_jump(
                jump_to,
                model_destination=model_destination,
                end_destination=end_destination,
            )

        # 2. Exit condition: A structured response was generated
        if "structured_response" in state:
            return end_destination

        # 3. Default: Continue the loop, there may have been an issue with structured
        # output generation, so we need to retry
        return model_destination

    return model_to_model


def _make_tools_to_model_edge(
    *,
    tool_node: ToolNode,
    model_destination: str,
    structured_output_tools: dict[str, OutputToolBinding[Any]],
    end_destination: str,
) -> Callable[[dict[str, Any]], str | None]:
    """OpenCode Style 的 tools → model 条件边。

    总是返回 model，不支持 return_direct 退出。
    让 LLM 处理工具结果后决定是否结束。
    """

    def tools_to_model(state: dict[str, Any]) -> str | None:
        # 注意：return_direct 和结构化输出工具不会直接退出
        # LLM 会在 model_node 中通过 finish_reason 决定是否结束
        return model_destination

    return tools_to_model


def _add_middleware_edge(
    graph: StateGraph[
        AgentState[ResponseT], ContextT, _InputAgentState, _OutputAgentState[ResponseT]
    ],
    *,
    name: str,
    default_destination: str,
    model_destination: str,
    end_destination: str,
    can_jump_to: list[JumpTo] | None,
) -> None:
    """为 middleware 节点添加边，统一处理 jump_to 语义。"""
    if can_jump_to:

        def jump_edge(state: dict[str, Any]) -> str:
            return (
                _resolve_jump(
                    state.get("jump_to"),
                    model_destination=model_destination,
                    end_destination=end_destination,
                )
                or default_destination
            )

        destinations = [default_destination]

        if "end" in can_jump_to:
            destinations.append(end_destination)
        if "tools" in can_jump_to:
            destinations.append("tools")
        if "model" in can_jump_to and name != model_destination:
            destinations.append(model_destination)

        graph.add_conditional_edges(name, RunnableCallable(jump_edge, trace=False), destinations)

    else:
        graph.add_edge(name, default_destination)


__all__ = [
    "_make_model_to_tools_edge",
    "_make_model_to_model_edge",
    "_make_tools_to_model_edge",
    "_add_middleware_edge",
    "_make_loop_controller_edge",
]

