"""
OpenCode Style agent 的 middleware 组合与 Tool wrapper 逻辑。

本模块从 `factory.py` 中拆分出与业务无关的中间件基础设施：
- 将多个 `wrap_model_call` / `awrap_model_call` 处理器按顺序合并为一个
- 将多个 tool-call wrapper（同步 / 异步）合并为单一包装函数
- 读取 `AgentMiddleware` 上的 `can_jump_to` 配置，供图路由层使用

与 LangChain 原始实现相比，这里只是对 middleware 行为做轻量包装和组合，
保持完全兼容的同时，方便在不同 agent 工厂中复用该逻辑。
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Awaitable, Callable, Generic, List, Sequence

from langchain.agents.middleware.types import (
    AgentMiddleware,
    ContextT,
    ExtendedModelResponse,
    JumpTo,
    ModelRequest,
    ModelResponse,
    ResponseT,
    ToolCallRequest,
    ToolCallWrapper,
)
from langchain_core.messages import AIMessage, ToolMessage
from langgraph.types import Command


@dataclass
class _ComposedExtendedModelResponse(Generic[ResponseT]):
    """Internal result from composed ``wrap_model_call`` middleware.

    Unlike ``ExtendedModelResponse`` (user-facing, single command), this holds the
    full list of commands accumulated across all middleware layers during
    composition.
    """

    model_response: ModelResponse[ResponseT]
    """The underlying model response."""

    commands: List[Command[Any]] = field(default_factory=list)
    """Commands accumulated from all middleware layers (inner-first, then outer)."""


def _normalize_to_model_response(
    result: ModelResponse | AIMessage | ExtendedModelResponse,
) -> ModelResponse:
    """Normalize middleware return value to ModelResponse."""
    if isinstance(result, AIMessage):
        return ModelResponse(result=[result], structured_response=None)
    if isinstance(result, ExtendedModelResponse):
        return result.model_response
    return result


def _chain_model_call_handlers(
    handlers: Sequence[
        Callable[[ModelRequest[ContextT], Callable[[ModelRequest[ContextT]], ModelResponse]],
                 ModelResponse | AIMessage | ExtendedModelResponse]
    ],
) -> Callable[
    [ModelRequest[ContextT], Callable[[ModelRequest[ContextT]], ModelResponse]],
    _ComposedExtendedModelResponse,
] | None:
    """Compose multiple ``wrap_model_call`` handlers into single middleware stack."""
    if not handlers:
        return None

    def _to_composed_result(
        result: ModelResponse | AIMessage | ExtendedModelResponse | _ComposedExtendedModelResponse,
        extra_commands: List[Command[Any]] | None = None,
    ) -> _ComposedExtendedModelResponse:
        commands: List[Command[Any]] = list(extra_commands or [])
        if isinstance(result, _ComposedExtendedModelResponse):
            commands.extend(result.commands)
            model_response = result.model_response
        elif isinstance(result, ExtendedModelResponse):
            model_response = result.model_response
            if result.command is not None:
                commands.append(result.command)
        else:
            model_response = _normalize_to_model_response(result)

        return _ComposedExtendedModelResponse(model_response=model_response, commands=commands)

    if len(handlers) == 1:
        single_handler = handlers[0]

        def normalized_single(
            request: ModelRequest[ContextT],
            handler: Callable[[ModelRequest[ContextT]], ModelResponse],
        ) -> _ComposedExtendedModelResponse:
            return _to_composed_result(single_handler(request, handler))

        return normalized_single

    def compose_two(
        outer: Callable[
            [ModelRequest[ContextT], Callable[[ModelRequest[ContextT]], ModelResponse]],
            ModelResponse | AIMessage | ExtendedModelResponse | _ComposedExtendedModelResponse,
        ],
        inner: Callable[
            [ModelRequest[ContextT], Callable[[ModelRequest[ContextT]], ModelResponse]],
            ModelResponse | AIMessage | ExtendedModelResponse | _ComposedExtendedModelResponse,
        ],
    ) -> Callable[
        [ModelRequest[ContextT], Callable[[ModelRequest[ContextT]], ModelResponse]],
        _ComposedExtendedModelResponse,
    ]:
        """Compose two handlers where outer wraps inner."""

        def composed(
            request: ModelRequest[ContextT],
            handler: Callable[[ModelRequest[ContextT]], ModelResponse],
        ) -> _ComposedExtendedModelResponse:
            accumulated_commands: List[Command[Any]] = []

            def inner_handler(req: ModelRequest[ContextT]) -> ModelResponse:
                accumulated_commands.clear()
                inner_result = inner(req, handler)
                if isinstance(inner_result, _ComposedExtendedModelResponse):
                    accumulated_commands.extend(inner_result.commands)
                    return inner_result.model_response
                if isinstance(inner_result, ExtendedModelResponse):
                    if inner_result.command is not None:
                        accumulated_commands.append(inner_result.command)
                    return inner_result.model_response
                return _normalize_to_model_response(inner_result)

            outer_result = outer(request, inner_handler)
            return _to_composed_result(
                outer_result,
                extra_commands=accumulated_commands or None,
            )

        return composed

    composed_handler = compose_two(handlers[-2], handlers[-1])
    for h in reversed(handlers[:-2]):
        composed_handler = compose_two(h, composed_handler)

    return composed_handler


def _chain_async_model_call_handlers(
    handlers: Sequence[
        Callable[
            [ModelRequest[ContextT], Callable[[ModelRequest[ContextT]], Awaitable[ModelResponse]]],
            Awaitable[ModelResponse | AIMessage | ExtendedModelResponse],
        ]
    ],
) -> Callable[
    [ModelRequest[ContextT], Callable[[ModelRequest[ContextT]], Awaitable[ModelResponse]]],
    Awaitable[_ComposedExtendedModelResponse],
] | None:
    """Compose multiple async ``wrap_model_call`` handlers into single middleware stack."""
    if not handlers:
        return None

    def _to_composed_result(
        result: ModelResponse | AIMessage | ExtendedModelResponse | _ComposedExtendedModelResponse,
        extra_commands: List[Command[Any]] | None = None,
    ) -> _ComposedExtendedModelResponse:
        commands: List[Command[Any]] = list(extra_commands or [])
        if isinstance(result, _ComposedExtendedModelResponse):
            commands.extend(result.commands)
            model_response = result.model_response
        elif isinstance(result, ExtendedModelResponse):
            model_response = result.model_response
            if result.command is not None:
                commands.append(result.command)
        else:
            model_response = _normalize_to_model_response(result)

        return _ComposedExtendedModelResponse(model_response=model_response, commands=commands)

    if len(handlers) == 1:
        single_handler = handlers[0]

        async def normalized_single(
            request: ModelRequest[ContextT],
            handler: Callable[[ModelRequest[ContextT]], Awaitable[ModelResponse]],
        ) -> _ComposedExtendedModelResponse:
            return _to_composed_result(await single_handler(request, handler))

        return normalized_single

    def compose_two(
        outer: Callable[
            [ModelRequest[ContextT], Callable[[ModelRequest[ContextT]], Awaitable[ModelResponse]]],
            Awaitable[ModelResponse | AIMessage | ExtendedModelResponse | _ComposedExtendedModelResponse],
        ],
        inner: Callable[
            [ModelRequest[ContextT], Callable[[ModelRequest[ContextT]], Awaitable[ModelResponse]]],
            Awaitable[ModelResponse | AIMessage | ExtendedModelResponse | _ComposedExtendedModelResponse],
        ],
    ) -> Callable[
        [ModelRequest[ContextT], Callable[[ModelRequest[ContextT]], Awaitable[ModelResponse]]],
        Awaitable[_ComposedExtendedModelResponse],
    ]:
        """Compose two async handlers where outer wraps inner."""

        async def composed(
            request: ModelRequest[ContextT],
            handler: Callable[[ModelRequest[ContextT]], Awaitable[ModelResponse]],
        ) -> _ComposedExtendedModelResponse:
            accumulated_commands: List[Command[Any]] = []

            async def inner_handler(req: ModelRequest[ContextT]) -> ModelResponse:
                accumulated_commands.clear()
                inner_result = await inner(req, handler)
                if isinstance(inner_result, _ComposedExtendedModelResponse):
                    accumulated_commands.extend(inner_result.commands)
                    return inner_result.model_response
                if isinstance(inner_result, ExtendedModelResponse):
                    if inner_result.command is not None:
                        accumulated_commands.append(inner_result.command)
                    return inner_result.model_response
                return _normalize_to_model_response(inner_result)

            outer_result = await outer(request, inner_handler)
            return _to_composed_result(
                outer_result,
                extra_commands=accumulated_commands or None,
            )

        return composed

    composed_handler = compose_two(handlers[-2], handlers[-1])
    for h in reversed(handlers[:-2]):
        composed_handler = compose_two(h, composed_handler)

    return composed_handler


def _chain_tool_call_wrappers(
    wrappers: Sequence[ToolCallWrapper],
) -> ToolCallWrapper | None:
    """Compose wrappers into middleware stack (first = outermost)."""
    if not wrappers:
        return None

    if len(wrappers) == 1:
        return wrappers[0]

    def compose_two(outer: ToolCallWrapper, inner: ToolCallWrapper) -> ToolCallWrapper:
        """Compose two wrappers where outer wraps inner."""

        def composed(
            request: ToolCallRequest,
            execute: Callable[[ToolCallRequest], ToolMessage | Command[Any]],
        ) -> ToolMessage | Command[Any]:
            def call_inner(req: ToolCallRequest) -> ToolMessage | Command[Any]:
                return inner(req, execute)

            return outer(request, call_inner)

        return composed

    result = wrappers[-1]
    for wrapper in reversed(wrappers[:-1]):
        result = compose_two(wrapper, result)

    return result


def _chain_async_tool_call_wrappers(
    wrappers: Sequence[
        Callable[
            [ToolCallRequest, Callable[[ToolCallRequest], Awaitable[ToolMessage | Command[Any]]]],
            Awaitable[ToolMessage | Command[Any]],
        ]
    ],
) -> (
    Callable[
        [ToolCallRequest, Callable[[ToolCallRequest], Awaitable[ToolMessage | Command[Any]]]],
        Awaitable[ToolMessage | Command[Any]],
    ]
    | None
):
    """Compose async wrappers into middleware stack (first = outermost)."""
    if not wrappers:
        return None

    if len(wrappers) == 1:
        return wrappers[0]

    def compose_two(
        outer: Callable[
            [ToolCallRequest, Callable[[ToolCallRequest], Awaitable[ToolMessage | Command[Any]]]],
            Awaitable[ToolMessage | Command[Any]],
        ],
        inner: Callable[
            [ToolCallRequest, Callable[[ToolCallRequest], Awaitable[ToolMessage | Command[Any]]]],
            Awaitable[ToolMessage | Command[Any]],
        ],
    ) -> Callable[
        [ToolCallRequest, Callable[[ToolCallRequest], Awaitable[ToolMessage | Command[Any]]]],
        Awaitable[ToolMessage | Command[Any]],
    ]:
        """Compose two async wrappers where outer wraps inner."""

        async def composed(
            request: ToolCallRequest,
            execute: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command[Any]]],
        ) -> ToolMessage | Command[Any]:
            async def call_inner(req: ToolCallRequest) -> ToolMessage | Command[Any]:
                return await inner(req, execute)

            return await outer(request, call_inner)

        return composed

    result = wrappers[-1]
    for wrapper in reversed(wrappers[:-1]):
        result = compose_two(wrapper, result)

    return result


def _get_can_jump_to(middleware: AgentMiddleware[Any, Any], hook_name: str) -> list[JumpTo]:
    """Get the `can_jump_to` list from either sync or async hook methods."""
    base_sync_method = getattr(AgentMiddleware, hook_name, None)
    base_async_method = getattr(AgentMiddleware, f"a{hook_name}", None)

    sync_method = getattr(middleware.__class__, hook_name, None)
    if (
        sync_method
        and sync_method is not base_sync_method
        and hasattr(sync_method, "__can_jump_to__")
    ):
        return sync_method.__can_jump_to__

    async_method = getattr(middleware.__class__, f"a{hook_name}", None)
    if (
        async_method
        and async_method is not base_async_method
        and hasattr(async_method, "__can_jump_to__")
    ):
        return async_method.__can_jump_to__

    return []


__all__ = [
    "_ComposedExtendedModelResponse",
    "_chain_model_call_handlers",
    "_chain_async_model_call_handlers",
    "_chain_tool_call_wrappers",
    "_chain_async_tool_call_wrappers",
    "_get_can_jump_to",
]

