"""Redact credentials in MCP tool output before it reaches the model.

Covers every MCP route; in-process tools like ``Bash``/``Read`` are
guarded separately by ``agents.policy``.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import mcp.types as mt
from fastmcp.server.middleware import Middleware
from fastmcp.tools.base import ToolResult

from beaver_gateway.security.redact import redact, redact_data

if TYPE_CHECKING:
    from fastmcp.server.middleware import CallNext, MiddlewareContext

__all__ = ["RedactingMiddleware"]


def _redact_block(block: mt.ContentBlock) -> mt.ContentBlock:
    """Mask a content block's text; binary blocks pass through."""
    if isinstance(block, mt.TextContent):
        masked = redact(block.text)
        return (
            block if masked == block.text else block.model_copy(update={"text": masked})
        )
    if isinstance(block, mt.EmbeddedResource) and isinstance(
        block.resource, mt.TextResourceContents
    ):
        masked = redact(block.resource.text)
        if masked == block.resource.text:
            return block
        resource = block.resource.model_copy(update={"text": masked})
        return block.model_copy(update={"resource": resource})
    return block


class RedactingMiddleware(Middleware):
    """Mask credentials in whatever a tool call returns.

    Idempotent, so mounting a server both on its own namespace and in
    the ``/mcp/all`` bundle costs a second pass, not a double mask.
    """

    async def on_call_tool(
        self,
        context: MiddlewareContext[mt.CallToolRequestParams],
        call_next: CallNext[mt.CallToolRequestParams, ToolResult],
    ) -> ToolResult:
        """Mask the result.

        A non-``None`` ``meta`` on the masked copy takes the
        ``CallToolResult`` path that skips output-schema validation,
        since masked structured content may no longer match it.
        """
        result = await call_next(context)
        content = [_redact_block(block) for block in result.content]
        structured = redact_data(result.structured_content)
        if content == result.content and structured == result.structured_content:
            return result
        return ToolResult(
            content=content,
            structured_content=structured,
            meta=result.meta if result.meta is not None else {},
        )