"""Redact credentials on the way out of an MCP tool.

A tool result is a wider channel than the log. It goes straight into the
model's context, from there into the turn record in Postgres, and from
there into the markdown transcript in ``💬 чаты`` — which Obsidian Sync
carries off the machine. One ``komodo`` deploy returns the resolved
compose file, ``environment:`` block and all, so a single call can put
every secret of a stack into all four places at once.

One filter, not a list of exceptions: every MCP the model can reach is
built into a ``FastMCP`` by :mod:`beaver_gateway.mcp.internal_app` —
``python_tool`` bundles like komodo, stdio subprocesses, remote HTTP
servers — and every route into one of them runs its middleware chain.
That covers the per-namespace ``/mcp/<name>/`` mounts claude-code talks
to, the ``/mcp/all`` bundle, the external MCP frontend reverse-proxying
into both, and the Raycast backend's direct ``call_tool``. A new MCP in
``config.py`` is covered the day it is added, without anyone
remembering to list it here.

What this does not reach: tools that never touch a FastMCP server — the
gateway's own ``gateway`` tools, and everything claude-code runs inside
its own process (``Bash``, ``Read``). Those are guarded by
:mod:`beaver_gateway.agents.policy` and the vault mounts instead.
"""

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:
        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,
            # Masked structured content no longer has to satisfy the
            # tool's output schema; a non-None meta takes the
            # ``CallToolResult`` path that skips that validation, the
            # same trick ``ResponseLimitingMiddleware`` uses.
            meta=result.meta if result.meta is not None else {},
        )