"""
Pure logic for merging an upstream A2A agent card with LiteLLM-specific overrides.

The merge produces the card that LiteLLM exposes to A2A clients at
``/a2a/{agent_id}/.well-known/agent-card.json``. The upstream card is taken as
the base; specific fields are replaced so all traffic flows through the proxy
and uses LiteLLM auth.
"""

import re
from collections.abc import Mapping
from copy import deepcopy
from typing import Any, Final, Literal

SupportedA2AVersion = Literal["0.3", "1.0"]

# Protocol versions LiteLLM can serve to A2A clients. The admin pins one per agent;
# responses are normalized to it regardless of the upstream agent's own version.
SUPPORTED_A2A_PROTOCOL_VERSIONS: Final[tuple[SupportedA2AVersion, ...]] = ("0.3", "1.0")

# Default served version when the agent card does not pin one.
LITELLM_A2A_PROTOCOL_VERSION: Final = "1.0"


_PROTOCOL_VERSION_PATTERN: Final = re.compile(
    r"^(\d+\.\d+)(?:\.\d+(?:-[0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*)?(?:\+[0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*)?)?$"
)


def normalize_protocol_version(version: object) -> SupportedA2AVersion | None:
    """Map a raw ``protocolVersion`` value to the supported canonical major.minor version.

    Accepts the bare major.minor convention of the 1.0 spec (``"0.3"``, ``"1.0"``) and the
    full semver forms older SDKs emit (``"0.3.0"``, ``"1.0.1"``, including prerelease and
    build suffixes like ``"0.3.0-rc1"``). Malformed strings, versions outside the
    supported set, and non-strings yield ``None``.
    """
    if not isinstance(version, str):
        return None
    match: Final = _PROTOCOL_VERSION_PATTERN.match(version)
    if match is None:
        return None
    major_minor: Final = match.group(1)
    return next((supported for supported in SUPPORTED_A2A_PROTOCOL_VERSIONS if supported == major_minor), None)


def resolve_served_protocol_version(card: Mapping[str, Any] | None) -> str:
    """Return the validated protocol version an agent card pins, else the default."""
    normalized: Final = normalize_protocol_version(card.get("protocolVersion") if card else None)
    return normalized if normalized is not None else LITELLM_A2A_PROTOCOL_VERSION


# Security scheme exposed by the LiteLLM-fronted agent card. Always replaces
# whatever upstream advertised — the client must authenticate to the proxy,
# not the upstream agent.
LITELLM_SECURITY_SCHEMES: Final[dict[str, dict[str, Any]]] = {
    "LiteLLMKey": {
        "type": "http",
        "scheme": "bearer",
        "description": "LiteLLM virtual key",
    },
}

LITELLM_SECURITY_REQUIREMENTS: Final[list[dict[str, list[str]]]] = [{"LiteLLMKey": []}]

# Capabilities LiteLLM can faithfully proxy today. Anything not in this set is
# dropped during merge so we don't advertise behavior the proxy can't deliver.
#
# TODO: re-enable ``streaming`` once the A2A streaming endpoint at
#   ``POST /a2a/{agent_id}/message/stream`` is exercised end-to-end with
#   cost tracking + guardrails. It's wired in ``a2a_endpoints.py`` but not
#   yet covered by tests, so we keep it gated on the upstream advertising it.
# TODO: ``pushNotifications`` — proxy has no webhook plumbing yet.
# TODO: ``extendedAgentCard`` — no separate authenticated-extended-card
#   endpoint exposed by the proxy.
# TODO: ``extensions`` — protocol extensions aren't validated/forwarded yet.
_ALLOWED_CAPABILITY_KEYS: Final = {"streaming"}

# v1.0 AgentCard top-level fields. Anything else is stripped from the merged
# card as a defense against upstream drift. ``supportedInterfaces`` is kept
# verbatim per product spec even though it is not in the v1.0 schema — clients
# that expect it will find it; clients that don't will ignore it.
#
# ``additionalInterfaces`` is deliberately excluded: it advertises alternate
# upstream URLs (HTTP/JSONRPC/gRPC backends) that, if persisted and served,
# would let authenticated agent callers reach the backend directly and bypass
# the proxy's auth/budget/logging. The proxy publishes its own entrypoint via
# ``supportedInterfaces`` instead.
_ALLOWED_TOP_LEVEL_KEYS: Final = {
    "protocolVersion",
    "name",
    "description",
    "version",
    "capabilities",
    "defaultInputModes",
    "defaultOutputModes",
    "skills",
    "preferredTransport",
    "supportedInterfaces",
    "iconUrl",
    "provider",
    "documentationUrl",
    "securitySchemes",
    "security",
    "supportsAuthenticatedExtendedCard",
    "signatures",
    # ``url`` is retained on the stored card because the runtime A2A invocation
    # path (``a2a_endpoints.py``) reads ``agent.agent_card_params['url']`` to
    # locate the upstream backend. The public ``/.well-known/agent-card.json``
    # endpoint rewrites this field to the proxy URL before serving it to
    # clients, so retaining it here does not leak the upstream to A2A callers.
    "url",
}

_DEFAULT_SKILLS: Final[list[dict[str, Any]]] = [
    {
        "id": "chat",
        "name": "Chat",
        "description": "Conversational interaction with the agent.",
        "tags": ["chat"],
    }
]

_DEFAULT_MODES: Final[list[str]] = ["text"]

# Fallback ``version`` when the upstream card omits the field. The A2A v1.0
# schema requires ``version`` on every card, so without this default the
# merged card would fail validation on clients that ``model_validate`` it.
_DEFAULT_AGENT_VERSION: Final = "1.0.0"


def _filter_capabilities(upstream_capabilities: Any) -> dict[str, Any]:
    """Return a capabilities dict containing only allowlisted, truthy keys."""
    if not isinstance(upstream_capabilities, dict):
        return {}
    return {
        key: value for key, value in upstream_capabilities.items() if key in _ALLOWED_CAPABILITY_KEYS and bool(value)
    }


def _default_litellm_provider(proxy_base_url: str) -> dict[str, str]:
    return {"organization": "LiteLLM Proxy", "url": proxy_base_url}


def merge_agent_card(
    upstream_card: Mapping[str, Any] | None,
    *,
    proxy_url: str,
    proxy_base_url: str,
    name: str | None = None,
    description: str | None = None,
) -> dict[str, Any]:
    """
    Build the LiteLLM-fronted agent card.

    Args:
        upstream_card: Card returned by the upstream agent's well-known endpoint.
            May be ``None``/empty when the upstream did not expose one.
        proxy_url: Full URL clients should hit to invoke this agent through
            the proxy, e.g. ``https://proxy.example.com/a2a/<agent_id>``.
        proxy_base_url: Root URL of the LiteLLM proxy, used as a fallback when
            we synthesize a provider record.
        name: User-supplied agent name from the LiteLLM UI. Takes precedence
            over the upstream card's ``name``.
        description: User-supplied description from the LiteLLM UI. Takes
            precedence over the upstream card's ``description``.

    Returns:
        A dict suitable for serving as the proxy's agent card. Only keys in
        the v1.0 AgentCard schema (plus ``supportedInterfaces``) are emitted.
    """
    base: Final[dict[str, Any]] = deepcopy(dict(upstream_card)) if upstream_card else {}

    # Keep the upstream ``url`` on the stored card: the runtime A2A
    # invocation path reads it from ``agent_card_params`` to know where to
    # proxy requests. The public well-known endpoint rewrites this field
    # to the proxy URL before exposing the card to clients.

    served_version: Final = resolve_served_protocol_version(upstream_card)
    base["protocolVersion"] = served_version

    if name:
        base["name"] = name
    if description:
        base["description"] = description

    if not base.get("version"):
        base["version"] = _DEFAULT_AGENT_VERSION

    base["capabilities"] = _filter_capabilities(base.get("capabilities"))

    if not base.get("skills"):
        base["skills"] = deepcopy(_DEFAULT_SKILLS)
    if not base.get("defaultInputModes"):
        base["defaultInputModes"] = list(_DEFAULT_MODES)
    if not base.get("defaultOutputModes"):
        base["defaultOutputModes"] = list(_DEFAULT_MODES)

    if not base.get("provider"):
        base["provider"] = _default_litellm_provider(proxy_base_url)

    base["supportedInterfaces"] = [
        {
            "url": proxy_url,
            "protocolBinding": "JSONRPC",
            "protocolVersion": served_version,
        }
    ]

    base["securitySchemes"] = deepcopy(LITELLM_SECURITY_SCHEMES)
    # Use the standard A2A/OpenAPI ``security`` field for requirements, not
    # the non-standard ``securityRequirements`` alias. The upstream's own
    # ``security`` selector is overwritten here because the proxy enforces its
    # own scheme regardless of what upstream required.
    base["security"] = deepcopy(LITELLM_SECURITY_REQUIREMENTS)

    return {key: value for key, value in base.items() if key in _ALLOWED_TOP_LEVEL_KEYS}
