from __future__ import annotations

from collections.abc import Mapping
from typing import Any

from dbos import DBOS

from pydantic_ai import ToolsetTool
from pydantic_ai.durable_exec._toolset import (
    CallToolResult,
    DurableMCPToolset,
    toolset_for_unit,
    unwrap_recorded_tool_call_result,
    wrap_tool_call_result,
)
from pydantic_ai.mcp import MCPToolset, ToolResult
from pydantic_ai.tools import AgentDepsT, RunContext, ToolDefinition

from ._utils import StepConfig, guard_enqueue_in_workflow


def dbosify_mcp_toolset(
    wrapped: MCPToolset[AgentDepsT], *, step_name_prefix: str, step_config: StepConfig
) -> DurableMCPToolset[AgentDepsT]:
    id_suffix = f'__{wrapped.id}' if wrapped.id else ''
    name = f'{step_name_prefix}__mcp_server{id_suffix}'

    @DBOS.step(name=f'{name}.get_tools', **(step_config or {}))
    async def get_tools_step(ctx: RunContext[AgentDepsT]) -> dict[str, ToolDefinition]:
        step_ctx = guard_enqueue_in_workflow(ctx)
        # The run holds one session for all of its steps; this is usually the step that opens it.
        async with toolset_for_unit(wrapped, step_ctx) as step_toolset:
            tools = await step_toolset.get_tools(step_ctx)
        return {tool_name: tool.tool_def for tool_name, tool in tools.items()}

    @DBOS.step(name=f'{name}.get_instructions', **(step_config or {}))
    async def get_instructions_step(ctx: RunContext[AgentDepsT]):
        step_ctx = guard_enqueue_in_workflow(ctx)
        # A server's instructions are captured when its session opens, so this step needs it open.
        async with toolset_for_unit(wrapped, step_ctx) as step_toolset:
            return await step_toolset.get_instructions(step_ctx)

    @DBOS.step(name=f'{name}.call_tool', **(step_config or {}))
    async def call_tool_step(
        tool_name: str,
        tool_args: dict[str, Any],
        ctx: RunContext[AgentDepsT],
        tool: ToolsetTool[AgentDepsT],
    ) -> CallToolResult:
        # The context is guarded because a `process_tool_call=` hook receives it and could enqueue.
        # DBOS has no selective non-retryable-exception support, so control-flow
        # exceptions must cross the step boundary as successful values.
        step_ctx = guard_enqueue_in_workflow(ctx)
        async with toolset_for_unit(wrapped, step_ctx) as step_toolset:
            return await wrap_tool_call_result(step_toolset.call_tool(tool_name, tool_args, step_ctx, tool))

    async def call_tool_operation(
        name: str,
        tool_args: dict[str, Any],
        *,
        ctx: RunContext[AgentDepsT],
        tool: ToolsetTool[AgentDepsT],
        config: Mapping[str, Any],
    ) -> ToolResult:
        # A recovering workflow may replay outputs this step recorded before it wrapped
        # control-flow exceptions as values; those recordings are the raw tool result.
        return unwrap_recorded_tool_call_result(await call_tool_step(name, tool_args, ctx, tool))

    return DurableMCPToolset(
        wrapped,
        # DBOS steps degrade gracefully to plain calls outside a workflow, so the durable
        # path is always taken — matching the previous DBOS wrapper, which never gated on
        # workflow state (outside a workflow, the step fallback still enters the server
        # for the run, around whichever step needs it first).
        in_durable_context=lambda: True,
        get_tools_operation=get_tools_step,
        get_instructions_operation=get_instructions_step,
        call_tool_operation=call_tool_operation,
        # DBOS takes no per-tool config; tool metadata is ignored, as before.
        resolve_tool_config=lambda tool, name: {},
        lifecycle='enter-in-durable-unit',
        durable_config=step_config,
    )


DBOSMCPToolset = DurableMCPToolset