"""MCP toolsets, rebuilt for every conversation turn.

The rule this enforces is the one thing the identity design depends on: **an MCP
client is built per turn, with tokens exchanged for that turn**, never held open
with a fixed token.

That is not a limitation of any framework — it is what keeps the blast radius of
a leaked token to one turn, and what makes `azp`/`sub` in the audit trail mean
something. PydanticAI takes `headers` at toolset construction, which is exactly
the shape this needs; the same constraint on other runtimes (langchain-mcp-adapters
issues #172/#194/#235, CrewAI's MCPServerAdapter) has the same answer.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import httpx
from pydantic_ai.mcp import MCPToolset, StreamableHttpTransport
from pydantic_ai.toolsets import (
    AbstractToolset,
    ApprovalRequiredToolset,
    FilteredToolset,
    PrefixedToolset,
)

from .approval import is_write, needs_approval
from .config import Settings
from .identity import TokenExchanger


@dataclass(frozen=True)
class McpTarget:
    """One MCP server this agent is allowed to reach."""

    name: str
    url: str
    audience_scope: str
    read_only: bool = False
    """Refuse this server's write tools inside the agent as well as at the gateway.

    Belt and braces on purpose. The gateway already refuses what the end user's
    roles do not allow, but an admin reaching this agent would otherwise be
    offered `create-order` — a tool this agent has no business calling whoever
    is asking.
    """


def targets(settings: Settings) -> list[McpTarget]:
    """The catalog domain read/write, plus sales read-only.

    sales is here to answer "does this product actually sell?" before a pricing
    or discontinuation decision. Read only — this agent never touches orders,
    and the scope it requests reflects that.
    """
    return [
        McpTarget("catalog", settings.mcp_catalog_url, settings.mcp_catalog_scope),
        McpTarget("sales", settings.mcp_sales_url, settings.mcp_sales_scope, read_only=True),
    ]


async def build_toolsets(
    settings: Settings,
    exchanger: TokenExchanger,
    subject_token: str,
) -> list[MCPToolset]:
    """One toolset per MCP server, each carrying its own audience-scoped token.

    Several narrow tokens rather than one broad one, deliberately: a token
    minted for `mcp-catalog` is refused by the sales gateway, so a prompt
    injection that talks the agent into reaching the wrong domain is refused at
    the endpoint rather than succeeding quietly.
    """
    def client_factory(**kwargs: Any) -> httpx.AsyncClient:
        # fastmcp builds this client itself and would otherwise take httpx's
        # 5s default connect timeout. See Settings.http_timeout for why that is
        # not survivable against `.local` names, and what the failure looks like
        # when it bites: a bare "Failed to initialize server session".
        kwargs["verify"] = settings.verify
        kwargs["timeout"] = httpx.Timeout(settings.http_timeout, connect=settings.http_timeout)
        return httpx.AsyncClient(**kwargs)

    toolsets: list[AbstractToolset[None]] = []
    for target in targets(settings):
        token = await exchanger.exchange(subject_token, target.audience_scope)
        toolset: AbstractToolset[None] = MCPToolset(
            StreamableHttpTransport(
                url=target.url,
                headers={"Authorization": f"Bearer {token}"},
                # No `verify=` here: the factory below sets it, and passing
                # both makes fastmcp ignore this one with a warning.
                httpx_client_factory=client_factory,
            ),
            id=target.name,
            init_timeout=settings.http_timeout,
        )

        if target.read_only:
            toolset = FilteredToolset(toolset, lambda ctx, tool: not is_write(tool))

        # Writes stop and wait for a person. The predicate reads the MCP
        # `readOnlyHint` the server publishes, so the server decides what counts
        # as a write and this agent keeps no list of tool names in step with it.
        toolset = ApprovalRequiredToolset(toolset, needs_approval)

        # Two servers today, more later. A collision between two domains'
        # namespaces would otherwise be silent — one tool shadowing another of
        # the same name and a different meaning.
        toolsets.append(PrefixedToolset(toolset, target.name))

    return toolsets