"""The v1 <-> v2 bridge for the credential resolver.

These edge functions translate v1's request objects into the resolver's typed inputs and map
its typed errors onto the proxy's public exception contract. They import v1 and live outside the
package's public surface so the resolver core (``resolver.py`` / ``types.py``) stays v1-free.
Nothing wires them into ``_create_mcp_client`` yet.

``to_server_spec`` maps only the modes the resolver has gone live for, returning ``None`` for
every other mode so the caller defers to v1 (parity-safe); it grows one branch per migrated mode.
"""

from __future__ import annotations

import base64
from typing import TYPE_CHECKING, Final, Literal, NoReturn

from fastapi import HTTPException
from pydantic import SecretStr
from typing_extensions import assert_never

from litellm.proxy._experimental.mcp_server.oauth_utils import resolve_upstream_resource
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
    DEFAULT_CREDENTIAL_HEADER,
    ApiKeyConfig,
    AuthorizationCodeConfig,
    ClientAuth,
    ClientCredentialsConfig,
    ClientSecretAuth,
    CredError,
    IdJagConfig,
    NoneConfig,
    PassthroughConfig,
    PrivateKeyJwtAuth,
    ServerSpec,
    SharedKey,
    Subject,
    TokenExchangeConfig,
)
from litellm.types.mcp import DEFAULT_SUBJECT_TOKEN_TYPE, MCPAuth

if TYPE_CHECKING:
    from litellm.proxy._types import UserAPIKeyAuth
    from litellm.types.mcp_server.mcp_server_manager import MCPServer

_TOKEN_EXCHANGE_SUBJECT_TOKEN_DEFAULT: Final = "urn:ietf:params:oauth:token-type:access_token"
_ID_JAG_SUBJECT_TOKEN_DEFAULT: Final = "urn:ietf:params:oauth:token-type:id_token"


def token_header(server: MCPServer) -> str:
    """The upstream header this server's resolved credential occupies.

    One owner for every arm, so no spec builder spells the default itself and a server can never
    hand two arms different answers.
    """
    return server.upstream_token_header or DEFAULT_CREDENTIAL_HEADER


def to_subject(user_api_key_auth: UserAPIKeyAuth | None, subject_token: str | None) -> Subject:
    """Map v1's authenticated principal onto the resolver's Subject.

    tenant_id / subject_id are empty for an unauthenticated caller; the per-user arms must reject
    an empty subject rather than share one credential slot across callers.
    """
    inbound: Final = SecretStr(subject_token) if subject_token else None
    if user_api_key_auth is None:
        return Subject(tenant_id="", subject_id="", inbound_token=inbound)
    return Subject(
        tenant_id=user_api_key_auth.org_id or user_api_key_auth.team_id or "",
        subject_id=user_api_key_auth.user_id or "",
        inbound_token=inbound,
    )


def to_server_spec(server: MCPServer) -> ServerSpec | None:
    """Map a v1 server onto a ServerSpec for a migrated mode, or None to defer to v1.

    BYOK is the per-user source of the ``api_key`` mode; its scheme rides on ``auth_type`` just
    like a shared key, but the value is per-user and not migrated yet, so a BYOK server defers
    to v1 regardless of ``auth_type`` (this guard is the seam the BYOK arm replaces later).

    Dispatches on the declared ``auth_type``. The match is exhaustive over ``MCPAuthType`` with
    an ``assert_never`` tail, so a newly added auth mode fails the type gate here until it is
    explicitly mapped or explicitly deferred, rather than silently falling through to v1. Live
    modes: ``none``, the static-header family (``api_key`` plus the Authorization schemes,
    all shared-key), ``oauth2`` per-user tokens (``authorization_code``), ``oauth2`` M2M
    (``client_credentials``), ``oauth2_token_exchange`` (OBO), and the client-forwarded token
    modes ``true_passthrough`` / ``oauth_delegate`` (``PassthroughConfig``); delegated/passthrough
    oauth2 and SigV4 return None and stay on v1.
    """
    if server.is_byok:
        return None  # per-user BYOK source not migrated yet -> defer to v1 (any auth_type)
    resource: Final = server.url or server.server_id
    auth_type: Final = server.auth_type
    match auth_type:
        case None | MCPAuth.none:
            if server.is_oauth_passthrough:
                return None  # passthrough is not migrated yet -> defer to v1
            return ServerSpec(server_id=server.server_id, resource=resource, config=NoneConfig())
        case MCPAuth.api_key:
            return _shared_key_spec(server, resource, "X-API-Key", "")
        case MCPAuth.bearer_token:
            return _shared_key_spec(server, resource, "Authorization", "Bearer")
        case MCPAuth.token:
            return _shared_key_spec(server, resource, "Authorization", "token")
        case MCPAuth.authorization:
            return _shared_key_spec(server, resource, "Authorization", "")
        case MCPAuth.basic:
            return _shared_key_spec(server, resource, "Authorization", "Basic", encode=True)
        case MCPAuth.oauth2:
            return _oauth2_spec(server, resource)
        case MCPAuth.oauth2_id_jag:
            return _id_jag_spec(server, resource)
        case MCPAuth.true_passthrough | MCPAuth.oauth_delegate:
            return ServerSpec(server_id=server.server_id, resource=resource, config=PassthroughConfig())
        case MCPAuth.oauth2_token_exchange:
            return _token_exchange_spec(server, resource)
        case MCPAuth.aws_sigv4:
            return None  # SigV4 is not migrated yet -> defer to v1
    assert_never(auth_type)


def _oauth2_spec(server: MCPServer, resource: str) -> ServerSpec | None:
    """Dispatch the oauth2 auth_type across its sub-modes: M2M, gateway-managed interactive, or v1.

    ``client_credentials`` (the explicit ``oauth2_flow`` opt-in) builds the M2M spec, per-user
    ``authorization_code`` without upstream delegation builds the interactive spec, and the
    delegate/passthrough shapes defer to v1 (None).
    """
    if server.has_client_credentials:
        return _client_credentials_spec(server, resource)
    if server.needs_user_oauth_token and not server.delegate_auth_to_upstream:
        return ServerSpec(
            server_id=server.server_id,
            resource=resource,
            config=AuthorizationCodeConfig(header_name=token_header(server)),
        )
    return None


def _client_credentials_spec(server: MCPServer, resource: str) -> ServerSpec:
    """Build a client_credentials (M2M) spec; the explicit ``oauth2_flow`` opt-in owns the server.

    Missing grant fields (``client_id``/``client_secret``/``token_url``) are NOT a reason to defer:
    v1 would connect unauthenticated and the upstream's 401 gets absorbed into an empty tool list,
    so the arm fails closed with ``misconfigured`` instead, naming the missing fields (mirrors the
    OBO ownership rule). ``audience`` is forwarded only when the operator set it; a missing one is
    omitted, not derived, since a fabricated value risks the IdP rejecting the grant.
    """
    return ServerSpec(
        server_id=server.server_id,
        resource=resource,
        config=ClientCredentialsConfig(
            header_name=token_header(server),
            client_id=server.client_id,
            client_secret=SecretStr(server.client_secret) if server.client_secret else None,
            token_url=server.effective_token_url,
            scopes=tuple(server.scopes or ()),
            audience=server.audience,
            upstream_resource=resolve_upstream_resource(server),
            token_endpoint_auth_method=server.token_endpoint_auth_method,
        ),
    )


def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec | None:
    """Build a token_exchange (OBO) spec, or defer (None) when it is not OBO-configured.

    An OBO server with ``client_id``/``client_secret`` is owned by the v2 arm even if the
    ``token_exchange_endpoint``/``token_url`` is absent: a missing endpoint then fails closed (412) at
    the exchanger rather than silently deferring to v1 and connecting unauthenticated, since the
    gateway must not guess the IdP or fall back to a weaker source. Without client credentials there is
    nothing to own, so the server stays on v1 (parity-safe). ``profile`` selects the wire dialect
    (``rfc8693`` default, ``entra_obo`` for Microsoft Entra On-Behalf-Of); an unrecognized value
    normalizes to ``rfc8693`` so a bad config value cannot crash spec-building. ``audience`` is
    forwarded only when the operator set it; a missing one is omitted, not derived.
    """
    endpoint: Final = server.token_exchange_endpoint or server.effective_token_url
    if not server.client_id or not server.client_secret:
        return None
    profile: Final[Literal["rfc8693", "entra_obo"]] = (
        "entra_obo" if server.token_exchange_profile == "entra_obo" else "rfc8693"
    )
    return ServerSpec(
        server_id=server.server_id,
        resource=resource,
        config=TokenExchangeConfig(
            header_name=token_header(server),
            profile=profile,
            subject_token_type=server.subject_token_type or DEFAULT_SUBJECT_TOKEN_TYPE,
            token_exchange_endpoint=endpoint,
            audience=server.audience,
            client_id=server.client_id,
            client_secret=SecretStr(server.client_secret),
            token_endpoint_auth_method=server.token_endpoint_auth_method,
            scopes=tuple(server.scopes or ()),
        ),
    )


def _shared_key_spec(
    server: MCPServer,
    resource: str,
    header_name: str,
    value_prefix: str,
    *,
    encode: bool = False,
) -> ServerSpec | None:
    """Build an api_key spec from the server's static token, or defer (None) if it is absent.

    Covers the whole shared-key static-header family: ``api_key`` on ``X-API-Key`` and the
    Authorization schemes (bearer / token / authorization sent verbatim, basic base64-encoded).
    """
    token: Final = server.authentication_token
    if not token:
        return None  # no key configured -> defer to v1 (parity-safe)
    value: Final = base64.b64encode(token.encode("utf-8")).decode() if encode else token
    return ServerSpec(
        server_id=server.server_id,
        resource=resource,
        config=ApiKeyConfig(
            header_name=server.upstream_token_header or header_name,
            value_prefix=value_prefix,
            key_source=SharedKey(value=SecretStr(value)),
        ),
    )


def _id_jag_spec(server: MCPServer, resource: str) -> ServerSpec | None:
    """Build an ID-JAG spec from the v1 server's raw fields, or defer (None) if half-configured.

    The enum already routes here, but a server missing an endpoint, ``client_id``, or any client-auth
    secret would make ``IdJagConfig`` raise at construction; returning None instead defers to v1 so a
    partially configured server does not 500. ``token_exchange_endpoint`` is leg 1 (the IdP org AS);
    leg 2 is ``id_jag_resource_token_endpoint`` (the upstream resource AS).
    """
    org_token_endpoint: Final = server.token_exchange_endpoint
    resource_token_endpoint: Final = server.id_jag_resource_token_endpoint
    client_id: Final = server.client_id
    client_auth: Final = _id_jag_client_auth(server)
    if not org_token_endpoint or not resource_token_endpoint or not client_id or client_auth is None:
        return None
    return ServerSpec(
        server_id=server.server_id,
        resource=resource,
        config=IdJagConfig(
            header_name=token_header(server),
            org_token_endpoint=org_token_endpoint,
            resource_token_endpoint=resource_token_endpoint,
            client_id=client_id,
            client_auth=client_auth,
            subject_token_type=_id_jag_subject_token_type(server),
            audience=server.audience,
            resource=server.id_jag_resource,
            scopes=tuple(server.scopes or ()),
        ),
    )


def _id_jag_client_auth(server: MCPServer) -> ClientAuth | None:
    """Private-key JWT when a key is configured, else client_secret, else None (defer to v1)."""
    if server.client_private_key:
        return PrivateKeyJwtAuth(
            private_key=SecretStr(server.client_private_key),
            key_id=server.client_private_key_id,
            signing_alg=server.client_assertion_signing_alg,
        )
    if server.client_secret:
        return ClientSecretAuth(client_secret=SecretStr(server.client_secret))
    return None


def _id_jag_subject_token_type(server: MCPServer) -> str:
    """ID-JAG asserts the user's id_token, so the token-exchange access_token default maps to id_token;
    an explicitly configured value (e.g. a SAML2 assertion type) is honored verbatim."""
    configured: Final = server.subject_token_type
    if configured and configured != _TOKEN_EXCHANGE_SUBJECT_TOKEN_DEFAULT:
        return configured
    return _ID_JAG_SUBJECT_TOKEN_DEFAULT


def raise_public(error: CredError) -> NoReturn:
    """Map a resolver CredError onto the proxy's public HTTP contract. The one edge that raises."""
    match error.tag:
        case "unauthorized":
            challenge: Final = error.unauthorized
            raise HTTPException(
                status_code=401,
                detail=challenge.body if challenge.body is not None else error.summary,
                headers=({"WWW-Authenticate": challenge.www_authenticate} if challenge.www_authenticate else None),
            )
        case "misconfigured":
            raise HTTPException(status_code=500, detail=error.summary)
        case "upstream_unavailable":
            raise HTTPException(status_code=503, detail=error.summary)
        case "unsupported_mode":
            raise HTTPException(status_code=500, detail=error.summary)
        case "precondition_required":
            raise HTTPException(status_code=412, detail=error.summary)
        case "not_implemented":
            raise HTTPException(status_code=501, detail=error.summary)
    assert_never(error.tag)


def oauth_protected_resource_path(root_path: str, server: MCPServer) -> str:
    """The server's RFC 9728 Protected Resource Metadata path, the shared anchor of both challenges.

    ``root_path`` is the proxy's ``SERVER_ROOT_PATH``, resolved by the caller (the imperative shell)
    so this stays a pure function of its inputs; ``"/"`` and ``""`` both mean no prefix. The path is
    relative, so it resolves against the caller's own host (correct even behind a reverse proxy).
    """
    prefix: Final = "" if root_path == "/" else root_path
    name: Final = server.alias or server.server_name or server.name or server.server_id
    return f"/.well-known/oauth-protected-resource{prefix}/mcp/{name}"


def raise_user_oauth_challenge(server: MCPServer, *, root_path: str) -> NoReturn:
    """Raise the 401 an ``authorization_code`` server returns at egress when the user has no token.

    Points at the server's RFC 9728 Protected Resource Metadata, which names the upstream
    authorization server the client must complete OAuth with. The listing-phase 401 still emits the
    RFC 8414 ``authorization_uri`` form pending the format unification; both target the same server,
    so the difference is cosmetic.
    """
    resource_metadata: Final = oauth_protected_resource_path(root_path, server)
    raise HTTPException(
        status_code=401,
        detail="Unauthorized",
        headers={"WWW-Authenticate": f'Bearer resource_metadata="{resource_metadata}"'},
    )


def raise_token_exchange_challenge(
    server: MCPServer,
    *,
    root_path: str,
    claims: str | None = None,
) -> NoReturn:
    """Raise the RFC 9728 / RFC 6750 challenge an OBO (``token_exchange``) server returns when the
    caller's subject token is missing or the IdP rejected it.

    Points at the server's Protected Resource Metadata, whose ``authorization_servers`` names the IdP
    the client must SSO with to obtain a subject token; ``error="invalid_token"`` tells a
    spec-compliant MCP client to discover that AS and retry with a fresh bearer. Mirrors
    ``raise_user_oauth_challenge`` but for the exchange flow: there is no gateway-side browser OAuth —
    the client re-authenticates directly with the IdP, and LiteLLM then exchanges the resulting token.

    An IdP step-up rejection (Entra Conditional Access / CAE) passes its ``claims`` blob. Per the
    Microsoft claims-challenge format the challenge then uses ``error="insufficient_claims"`` (the
    value MSAL-family clients key on) and carries the claims base64-encoded in a ``claims`` parameter
    the client replays to the IdP to satisfy the step-up. Without a claims blob the challenge keeps
    ``error="invalid_token"`` and is byte-identical to the static one. Both the error value (one of
    two literals) and the base64 claims draw from a fixed alphabet, so nothing from the IdP body
    reaches the header unescaped.
    """
    resource_metadata: Final = oauth_protected_resource_path(root_path, server)
    encoded_claims: Final = base64.b64encode(claims.encode()).decode() if claims else None
    error: Final = "insufficient_claims" if encoded_claims else "invalid_token"
    error_description: Final = (
        "Step-up authentication required; satisfy the returned claims challenge with the IdP and retry"
        if encoded_claims
        else "Missing or invalid subject token; authenticate with the IdP and retry"
    )
    www_authenticate: Final = ", ".join(
        (
            f'Bearer resource_metadata="{resource_metadata}"',
            f'error="{error}"',
            f'error_description="{error_description}"',
            *((f'claims="{encoded_claims}"',) if encoded_claims else ()),
        )
    )
    raise HTTPException(
        status_code=401,
        detail="Unauthorized",
        headers={"WWW-Authenticate": www_authenticate},
    )