"""Inject _meta into outgoing MCP tool calls.

Per docs/INTERFACES.md Section E, anima injects _meta into every outgoing
MCP tool call's arguments. This module provides the injection callback
that wraps MCP tool calls.

The PydanticAI process_tool_call hook allows customizing tool calls
including adding extra request metadata.
"""

# @shell_orchestration: Meta injection module orchestrates MCP callback attachment
# @invar:allow shell_result: Returns ToolResult from PydanticAI MCP

from collections.abc import Awaitable
from typing import TYPE_CHECKING, Any, Callable

from pydantic_ai.mcp import CallToolFunc, ToolResult
from pydantic_ai.tools import RunContext

from anima.core.schemas import AgentStateData

if TYPE_CHECKING:
    pass


# @invar:allow shell_result: Pure dict builder returns dict, not Result
# @shell_orchestration: Pure dict builder for _meta injection; retained in shell for MCP callback locality
def build_call_meta(state: AgentStateData) -> dict[str, Any]:
    """Build _meta dict for outgoing MCP tool calls.

    Per docs/INTERFACES.md Section E and contracts/meta.schema.json, _meta contains:
    - trace_id: Always present (end-to-end trace, required)
    - instance_id: Always present (agent instance)
    - persona_id: Only when PersonaSpec has it
    - event_id: Not set by anima (set by nervus dispatch)
    - idempotency_key: Not set by anima (set by tela per tool call)

    Note: depth and parent_id are internal AgentStateData fields for sub-agent
    nesting control and are NOT part of the _meta wire contract.

    Args:
        state: AgentStateData containing agent metadata.

    Returns:
        Dict with _meta fields for outgoing tool calls.

    >>> from anima.core.schemas import AgentStateData
    >>> state = AgentStateData(instance_id="inst_abc", persona_name="test", trace_id="tr_xyz")
    >>> meta = build_call_meta(state)
    >>> meta["trace_id"]
    'tr_xyz'
    >>> meta["instance_id"]
    'inst_abc'
    >>> sorted(meta.keys())
    ['instance_id', 'trace_id']

    >>> # With persona_id source from state.spec_id
    >>> state2 = AgentStateData(
    ...     instance_id="inst_uvw",
    ...     persona_name="test",
    ...     trace_id="tr_456",
    ...     spec_id="persona@abc123"
    ... )
    >>> meta2 = build_call_meta(state2)
    >>> meta2["persona_id"]
    'persona@abc123'
    >>> "spec_id" in meta2
    False

    >>> # spec_digest is not part of Meta and is omitted
    >>> state3 = AgentStateData(
    ...     instance_id="inst_full",
    ...     persona_name="test",
    ...     trace_id="tr_full",
    ...     spec_id="persona@v1",
    ...     spec_digest="sha256:abc123"
    ... )
    >>> meta3 = build_call_meta(state3)
    >>> sorted(meta3.keys())
    ['instance_id', 'persona_id', 'trace_id']

    >>> # depth and parent_id are NOT included in _meta (internal use only)
    >>> state4 = AgentStateData(
    ...     instance_id="inst_deep",
    ...     persona_name="test",
    ...     trace_id="tr_789",
    ...     spec_id="worker@v1",
    ...     spec_digest="sha256:xyz789",
    ...     depth=1,
    ...     parent_id="inst_parent"
    ... )
    >>> meta4 = build_call_meta(state4)
    >>> sorted(meta4.keys())
    ['instance_id', 'persona_id', 'trace_id']
    >>> "depth" in meta4
    False
    >>> "parent_id" in meta4
    False
    """
    # Per docs/INTERFACES.md Section E field presence rules:
    # - trace_id: Always (required by schema)
    # - instance_id: Always
    # - persona_id: Only if present
    # - event_id: Not set by anima (nervus sets it)
    # - idempotency_key: Not set by anima (tela sets it per tool call)
    #
    # Note: depth and parent_id are internal fields, NOT part of _meta contract
    meta: dict[str, Any] = {
        "trace_id": state.trace_id,
        "instance_id": state.instance_id,
    }

    # Optional fields from PersonaSpec
    if state.spec_id is not None:
        meta["persona_id"] = state.spec_id

    return meta


# @invar:allow shell_result: Returns ToolResult from async callback
# @shell_orchestration: Factory function for MCP callback creation; retained in shell for async closure locality
def create_meta_injection_callback(
    state: AgentStateData,
) -> Callable[
    [RunContext[Any], CallToolFunc, str, dict[str, Any]],
    Awaitable[ToolResult],
]:
    """Create a process_tool_call callback that injects _meta.

    This creates a callback suitable for passing to MCPServerStdio/MCPServerSSE
    via the `process_tool_call` parameter. The callback injects the _meta
    dict into tool arguments before forwarding to the underlying tool.

    Args:
        state: AgentStateData containing agent metadata for _meta injection.

    Returns:
        Async callback that wraps tool calls with _meta injection.

    Example usage:
        >>> from anima.core.schemas import AgentStateData
        >>> state = AgentStateData(
        ...     instance_id="inst_123",
        ...     persona_name="test",
        ...     trace_id="tr_abc",
        ...     depth=0
        ... )
        >>> callback = create_meta_injection_callback(state)
        >>> callback.__name__
        'inject_meta_callback'
    """

    async def inject_meta_callback(
        ctx: RunContext[Any],
        call_tool: CallToolFunc,
        name: str,
        args: dict[str, Any],
    ) -> ToolResult:
        """Inject _meta into outgoing MCP tool call.

        This callback:
        1. Builds _meta dict from agent state
        2. Passes _meta as metadata to the tool call
        3. Calls the underlying tool

        Per PydanticAI MCP implementation, the metadata parameter is passed
        directly to the MCP CallToolRequestParams._meta field.

        Args:
            ctx: RunContext from PydanticAI agent run.
            call_tool: The original tool call function (direct_call_tool).
            name: Name of the tool being called.
            args: Arguments to pass to the tool.

        Returns:
            ToolResult from the underlying tool call.
        """
        # Build _meta from agent state
        _meta = build_call_meta(state)

        # Pass _meta as metadata to the tool call
        # Per docs/INTERFACES.md Section E, this becomes the request metadata
        # Per PydanticAI MCP, the metadata param maps to CallToolRequestParams._meta
        return await call_tool(name, args, _meta)

    return inject_meta_callback


# @invar:allow shell_result: Returns MCP servers directly (no Result needed)
# @invar:allow dead_export: Wrapper function retained for future CLI integration; enables runtime meta injection config
# @shell_orchestration: MCP server wrapper orchestrating callback attachment; shell module for MCP integration
def wrap_mcp_servers_with_meta(
    mcp_servers: list[Any],
    state: AgentStateData,
) -> list[Any]:
    """Wrap MCP servers with _meta injection callback.

    Per docs/INTERFACES.md Section E, anima must inject _meta into every
    outgoing MCP tool call. This function wraps MCPServerStdio/MCPServerSSE
    instances with a process_tool_call callback that injects _meta.

    Args:
        mcp_servers: List of MCP server instances (MCPServerStdio/MCPServerSSE).
        state: AgentStateData containing agent metadata for _meta injection.

    Returns:
        List of MCP servers with _meta injection configured.

    Note:
        This function modifies servers in-place by setting process_tool_call
        if not already set. Servers with existing process_tool_call are not
        modified (preserving existing behavior).

    >>> from pydantic_ai.mcp import MCPServerStdio
    >>> from anima.core.schemas import AgentStateData
    >>> state = AgentStateData(instance_id="inst_1", persona_name="p", trace_id="tr_1")
    >>> server = MCPServerStdio("echo", [])
    >>> servers = wrap_mcp_servers_with_meta([server], state)
    >>> servers[0].process_tool_call is not None
    True
    """
    callback = create_meta_injection_callback(state)

    for server in mcp_servers:
        # Only set callback if not already configured
        # This preserves existing process_tool_call if present
        if server.process_tool_call is None:
            server.process_tool_call = callback

    return mcp_servers
