"""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 .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


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),
    ]


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 fails at
    Kong 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[MCPToolset] = []
    for target in targets(settings):
        token = await exchanger.exchange(subject_token, target.audience_scope)
        toolsets.append(
            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,
            )
        )
    return toolsets