import time
from typing import Any

from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext

from ..observability import (
    clear_context,
    log_tool_complete,
    log_tool_error,
    log_tool_start,
    set_tool_context,
)


class ToolTracingMiddleware(Middleware):
    """Middleware that traces tool execution with timing and MCP protocol logging.

    Sends tool execution metrics to the client via MCP protocol logging,
    and writes structured events to local JSONL log when MCP_LOGGING=safe|full.
    """

    async def on_call_tool(
        self,
        context: MiddlewareContext[Any],
        call_next: CallNext[Any, Any],
    ) -> Any:
        tool_name = getattr(context.message, "name", "unknown")
        params = getattr(context.message, "arguments", None)

        set_tool_context(tool_name)
        log_tool_start(tool_name, params)

        start = time.perf_counter()
        try:
            result = await call_next(context)
            duration_ms = (time.perf_counter() - start) * 1000

            result_keys = list(result.keys()) if isinstance(result, dict) else None
            log_tool_complete(tool_name, duration_ms, result_keys)
            await self._log_to_client(context, tool_name, duration_ms, success=True)

            return result
        except Exception as exc:
            duration_ms = (time.perf_counter() - start) * 1000
            log_tool_error(tool_name, duration_ms, str(exc), type(exc).__name__)
            await self._log_to_client(
                context, tool_name, duration_ms, success=False, error=str(exc)
            )
            raise
        finally:
            clear_context()

    async def _log_to_client(
        self,
        context: MiddlewareContext[Any],
        tool_name: str,
        duration_ms: float,
        *,
        success: bool,
        error: str | None = None,
    ) -> None:
        try:
            ctx = context.fastmcp_context
            if ctx is None:
                return

            status = "ok" if success else f"error: {error}"
            message = f"[{tool_name}] {status} ({duration_ms:.0f}ms)"

            if hasattr(ctx, "debug"):
                await ctx.debug(message)
        except Exception:  # nosec B110
            pass
