"""MCP server instrumentation wrapper."""

from __future__ import annotations

import time
from typing import Any, Protocol, cast, runtime_checkable

from .session import AALSession


@runtime_checkable
class McpServerLike(Protocol):
    """Duck-typed MCP server interface."""

    _registeredTools: dict[str, Any]  # noqa: N815


def wrap_mcp_server(server: McpServerLike, session: AALSession) -> None:
    """Wrap all registered MCP tools to log actions."""
    for tool_name in list(server._registeredTools.keys()):
        original_handler = server._registeredTools[tool_name]
        server._registeredTools[tool_name] = _make_wrapper(tool_name, original_handler, session)


def _make_wrapper(name: str, handler: Any, session: AALSession) -> Any:
    """Create a wrapper that logs tool calls."""

    async def wrapper(*args: Any, **kwargs: Any) -> Any:
        start = time.perf_counter()
        try:
            result = await handler(*args, **kwargs)
            latency_ms = round((time.perf_counter() - start) * 1000)
            action: dict[str, Any] = {
                "action_type": "TOOL_CALL",
                "tool_name": name,
                "status": "SUCCESS",
                "latency_ms": latency_ms,
            }
            if len(args) == 1 and isinstance(args[0], dict):
                action["parameters"] = cast(dict[str, object], args[0])
            if isinstance(result, dict):
                action["response"] = cast(dict[str, object], result)
            session.log_action(action)  # type: ignore[arg-type]
            return result
        except Exception as err:
            latency_ms = round((time.perf_counter() - start) * 1000)
            action_err: dict[str, Any] = {
                "action_type": "TOOL_CALL",
                "tool_name": name,
                "response": {"error": str(err)},
                "status": "ERROR",
                "latency_ms": latency_ms,
            }
            if len(args) == 1 and isinstance(args[0], dict):
                action_err["parameters"] = cast(dict[str, object], args[0])
            session.log_action(action_err)  # type: ignore[arg-type]
            raise

    return wrapper
