"""Unit tests for `mcp.server.request_state`: codec, security policy, and default principal binding."""

import base64
import string
from collections.abc import Callable
from typing import Any, cast

import pytest
from inline_snapshot import snapshot

from mcp.server.auth.middleware.auth_context import auth_context_var
from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser, authorization_context
from mcp.server.auth.provider import AccessToken, principal_components
from mcp.server.context import ServerRequestContext
from mcp.server.request_state import (
    AESGCMRequestStateCodec,
    InvalidRequestState,
    RequestStateSecurity,
    authenticated_principal,
)

_TOKEN_PREFIX = "v1."
_KID_LEN = 4
_NONCE_LEN = 12
_GCM_TAG_LEN = 16
_BODY_FLOOR = _KID_LEN + _NONCE_LEN + _GCM_TAG_LEN
_B64URL_ALPHABET = set(string.ascii_letters + string.digits + "-_")

_KEY_A = b"a" * 32
_KEY_B = b"b" * 32
_KEY_OLD = b"o" * 32
_KEY_NEW = b"n" * 32

# Distinctive plaintext: opacity and log-secrecy assertions search for it.
_PAYLOAD = b"sentinel-plaintext-3f9c"
# `InvalidRequestState` messages are short log-only reason codes, never payload.
_REASON_CODE_MAX_LEN = 40


def _b64u_nopad(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).decode().rstrip("=")


def _decode_body(token: str) -> bytes:
    body = token.removeprefix(_TOKEN_PREFIX)
    return base64.urlsafe_b64decode(body + "=" * (-len(body) % 4))


def _flip_body_byte(token: str, index: int) -> str:
    raw = bytearray(_decode_body(token))
    raw[index] ^= 0xFF
    return _TOKEN_PREFIX + _b64u_nopad(bytes(raw))


def _flip_prefix_char(token: str) -> str:
    return "x" + token[1:]


def _flip_kid_byte(token: str) -> str:
    return _flip_body_byte(token, 0)


def _flip_nonce_byte(token: str) -> str:
    return _flip_body_byte(token, _KID_LEN)


def _flip_ciphertext_byte(token: str) -> str:
    return _flip_body_byte(token, _KID_LEN + _NONCE_LEN)


def _flip_tag_byte(token: str) -> str:
    return _flip_body_byte(token, -1)


def _inject_junk_chars(body: str) -> str:
    return body[:10] + "!@\n*" + body[10:]


def _append_newline(body: str) -> str:
    return body + "\n"


def _append_padding(body: str) -> str:
    return body + "=" * (-len(body) % 4 or 4)


def _bare_context() -> ServerRequestContext[Any, Any]:
    return ServerRequestContext(
        session=cast("Any", None),
        lifespan_context={},
        protocol_version="2026-07-28",
        method="tools/call",
    )


class _StaticCodec:
    """Minimal `RequestStateCodec` stand-in for policy tests; no real crypto."""

    def seal(self, payload: bytes) -> str:
        return payload.hex()

    def unseal(self, token: str) -> bytes:
        return bytes.fromhex(token)


# -- AESGCMRequestStateCodec --------------------------------------------------


@pytest.mark.parametrize(
    "payload",
    [
        pytest.param(b"", id="empty"),
        pytest.param(b"plain ascii state", id="ascii"),
        pytest.param("ünïcødé – 状態".encode(), id="multi-byte-utf8"),
        pytest.param(bytes(range(256)), id="raw-binary"),
        pytest.param(bytes(range(256)) * 256, id="64KiB"),
    ],
)
def test_seal_unseal_round_trips_any_payload(payload: bytes) -> None:
    """SDK-defined: the codec is byte-transparent, so any payload survives seal/unseal unchanged."""
    codec = AESGCMRequestStateCodec([_KEY_A])
    assert codec.unseal(codec.seal(payload)) == payload


def test_a_sealed_token_is_v1_plus_unpadded_b64url_over_kid_nonce_and_ciphertext() -> None:
    """SDK-defined token format: "v1." plus unpadded base64url over kid(4) || nonce(12) || ciphertext+tag."""
    token = AESGCMRequestStateCodec([_KEY_A]).seal(_PAYLOAD)
    assert token.startswith(_TOKEN_PREFIX)
    body = token.removeprefix(_TOKEN_PREFIX)
    assert "=" not in body
    assert set(body) <= _B64URL_ALPHABET
    assert len(_decode_body(token)) == _KID_LEN + _NONCE_LEN + len(_PAYLOAD) + _GCM_TAG_LEN


def test_two_seals_of_the_same_payload_produce_distinct_tokens_that_both_unseal() -> None:
    """SDK-defined: every seal draws a fresh nonce, so identical payloads yield distinct tokens that both verify."""
    codec = AESGCMRequestStateCodec([_KEY_A])
    first = codec.seal(_PAYLOAD)
    second = codec.seal(_PAYLOAD)
    assert first != second
    assert codec.unseal(first) == _PAYLOAD
    assert codec.unseal(second) == _PAYLOAD


@pytest.mark.parametrize(
    "corrupt",
    [
        pytest.param(_flip_prefix_char, id="prefix-char"),
        pytest.param(_flip_kid_byte, id="kid-byte"),
        pytest.param(_flip_nonce_byte, id="nonce-byte"),
        pytest.param(_flip_ciphertext_byte, id="ciphertext-byte"),
        pytest.param(_flip_tag_byte, id="tag-byte"),
    ],
)
def test_a_token_corrupted_in_any_region_is_rejected_without_echoing_the_payload(
    corrupt: Callable[[str], str],
) -> None:
    """Spec-mandated (basic/patterns/mrtr, server requirement 4): any corrupted token region is rejected."""
    codec = AESGCMRequestStateCodec([_KEY_A])
    token = codec.seal(_PAYLOAD)
    with pytest.raises(InvalidRequestState) as exc:
        codec.unseal(corrupt(token))
    message = str(exc.value)
    assert len(message) <= _REASON_CODE_MAX_LEN
    assert _PAYLOAD.decode() not in message


@pytest.mark.parametrize(
    "token",
    [
        pytest.param("", id="empty-string"),
        pytest.param(_b64u_nopad(b"\x00" * 64), id="missing-prefix"),
        pytest.param(_TOKEN_PREFIX + "!!!not-base64!!!", id="garbage-after-prefix"),
        pytest.param(_TOKEN_PREFIX + _b64u_nopad(b"\x00" * (_BODY_FLOOR - 1)), id="below-floor"),
    ],
)
def test_a_structurally_malformed_token_is_rejected(token: str) -> None:
    """Spec-mandated (basic/patterns/mrtr, server requirement 4): tokens this codec never minted fail."""
    with pytest.raises(InvalidRequestState):
        AESGCMRequestStateCodec([_KEY_A]).unseal(token)


def test_a_token_minted_under_a_key_outside_the_ring_is_rejected_as_unknown_key() -> None:
    """Spec-mandated (basic/patterns/mrtr, server requirement 4): a foreign-key token fails as "unknown key"."""
    token = AESGCMRequestStateCodec([_KEY_A]).seal(_PAYLOAD)
    with pytest.raises(InvalidRequestState) as exc:
        AESGCMRequestStateCodec([_KEY_B]).unseal(token)
    assert str(exc.value) == "unknown key"


@pytest.mark.parametrize(
    "ring",
    [
        pytest.param([_KEY_OLD, _KEY_NEW], id="rotation-phase-1"),
        pytest.param([_KEY_NEW, _KEY_OLD], id="rotation-phase-2"),
    ],
)
def test_a_token_minted_under_the_old_key_unseals_under_any_ring_containing_it(ring: list[bytes]) -> None:
    """SDK-defined rotation: every ring key verifies, so old-key state survives both rollout phases."""
    token = AESGCMRequestStateCodec([_KEY_OLD]).seal(_PAYLOAD)
    assert AESGCMRequestStateCodec(ring).unseal(token) == _PAYLOAD


def test_the_first_ring_key_mints_and_later_ring_keys_only_verify() -> None:
    """SDK-defined rotation: keys[0] is the minter, so [new, old] state verifies under [new] but not [old]."""
    token = AESGCMRequestStateCodec([_KEY_NEW, _KEY_OLD]).seal(_PAYLOAD)
    assert AESGCMRequestStateCodec([_KEY_NEW]).unseal(token) == _PAYLOAD
    with pytest.raises(InvalidRequestState):
        AESGCMRequestStateCodec([_KEY_OLD]).unseal(token)


def test_a_token_minted_under_a_retired_key_is_rejected() -> None:
    """Spec-mandated (basic/patterns/mrtr, server requirement 4): retired-key state fails verification."""
    token = AESGCMRequestStateCodec([_KEY_OLD]).seal(_PAYLOAD)
    with pytest.raises(InvalidRequestState):
        AESGCMRequestStateCodec([_KEY_NEW]).unseal(token)


def test_an_empty_key_ring_is_rejected_at_construction() -> None:
    """SDK-defined: an empty ring is a configuration error caught at construction."""
    with pytest.raises(ValueError) as exc:
        AESGCMRequestStateCodec([])
    assert str(exc.value) == snapshot("AESGCMRequestStateCodec requires at least one key")


def test_a_key_shorter_than_32_bytes_is_rejected_with_generation_guidance() -> None:
    """SDK-defined: keys must carry at least 32 bytes; the error includes generation guidance."""
    with pytest.raises(ValueError) as exc:
        AESGCMRequestStateCodec([b"k" * 31])
    assert str(exc.value) == snapshot(
        "request-state keys must be at least 32 bytes of secret randomness; keys[0] is 31 bytes. "
        'Generate one with: python -c "import secrets; print(secrets.token_hex(32))"'
    )


def test_a_duplicate_key_in_the_ring_is_rejected_at_construction() -> None:
    """SDK-defined: duplicate ring keys are a rotation mistake caught at construction."""
    with pytest.raises(ValueError) as exc:
        AESGCMRequestStateCodec([_KEY_A, _KEY_A])
    assert str(exc.value) == snapshot("keys[1] duplicates an earlier ring key")


def test_a_non_key_typed_ring_entry_is_rejected_naming_its_index_and_type() -> None:
    """SDK-defined: a non-key ring entry raises a TypeError naming its index and type, in codec and policy."""
    with pytest.raises(TypeError) as exc:
        AESGCMRequestStateCodec([_KEY_A, cast("Any", 32)])
    assert str(exc.value) == snapshot("request-state keys must be bytes, bytearray, or str; keys[1] is int")
    with pytest.raises(TypeError) as exc:
        RequestStateSecurity(keys=[cast("Any", 32)])
    assert str(exc.value) == snapshot("request-state keys must be bytes, bytearray, or str; keys[0] is int")


def test_a_mixed_ring_of_bytes_bytearray_and_str_entries_still_works() -> None:
    """SDK-defined: bytes, bytearray, and str keys interoperate in one ring."""
    codec = AESGCMRequestStateCodec([_KEY_A, bytearray(_KEY_B), "c" * 32])
    assert codec.unseal(codec.seal(_PAYLOAD)) == _PAYLOAD
    assert codec.unseal(AESGCMRequestStateCodec([bytearray(_KEY_B)]).seal(_PAYLOAD)) == _PAYLOAD
    assert codec.unseal(AESGCMRequestStateCodec(["c" * 32]).seal(_PAYLOAD)) == _PAYLOAD


def test_a_str_key_is_equivalent_to_its_utf8_bytes_form() -> None:
    """SDK-defined: a str key is utf-8 encoded, so it is the same ring key as its bytes spelling."""
    token = AESGCMRequestStateCodec(["k" * 32]).seal(_PAYLOAD)
    assert AESGCMRequestStateCodec([b"k" * 32]).unseal(token) == _PAYLOAD


def test_bytearray_key_material_is_copied_at_construction() -> None:
    """SDK-defined: key bytes are copied at construction; mutating the caller's bytearray later has no effect."""
    material = bytearray(b"m" * 32)
    codec = AESGCMRequestStateCodec([cast("Any", material)])
    minted_before_mutation = codec.seal(_PAYLOAD)
    material[:] = b"X" * 32
    assert codec.unseal(minted_before_mutation) == _PAYLOAD
    assert AESGCMRequestStateCodec([b"m" * 32]).unseal(codec.seal(_PAYLOAD)) == _PAYLOAD


def test_the_token_reveals_the_payload_neither_in_its_text_nor_its_decoded_bytes() -> None:
    """SDK-defined: the token is encrypted, not merely signed, so the plaintext appears nowhere in it."""
    token = AESGCMRequestStateCodec([_KEY_A]).seal(_PAYLOAD)
    assert _PAYLOAD.decode() not in token
    assert _b64u_nopad(_PAYLOAD) not in token
    assert _PAYLOAD.hex() not in token
    assert _PAYLOAD not in _decode_body(token)


def test_every_substitution_of_the_final_token_character_is_rejected() -> None:
    """Spec-mandated (basic/patterns/mrtr, server requirement 4): canonical decoding
    rejects every final-character substitution despite base64 don't-care padding bits."""
    codec = AESGCMRequestStateCodec([_KEY_A])
    body = codec.seal(_PAYLOAD).removeprefix(_TOKEN_PREFIX)
    substitutions = [c for c in sorted(_B64URL_ALPHABET) if c != body[-1]]
    assert len(substitutions) == 63
    for c in substitutions:
        with pytest.raises(InvalidRequestState):
            codec.unseal(_TOKEN_PREFIX + body[:-1] + c)


@pytest.mark.parametrize(
    "mangle",
    [
        pytest.param(_inject_junk_chars, id="junk-chars-injected"),
        pytest.param(_append_newline, id="newline-appended"),
        pytest.param(_append_padding, id="padding-appended"),
    ],
)
def test_a_non_canonical_token_body_is_rejected(mangle: Callable[[str], str]) -> None:
    """Spec-mandated (basic/patterns/mrtr, server requirement 4): lax-decoder aliases of a token are rejected."""
    codec = AESGCMRequestStateCodec([_KEY_A])
    body = codec.seal(_PAYLOAD).removeprefix(_TOKEN_PREFIX)
    with pytest.raises(InvalidRequestState):
        codec.unseal(_TOKEN_PREFIX + mangle(body))


def test_a_token_reprefixed_to_a_future_format_version_is_rejected() -> None:
    """Spec-mandated (basic/patterns/mrtr, server requirement 4): the prefix is tag-bound; "v2." replay fails."""
    codec = AESGCMRequestStateCodec([_KEY_A])
    token = codec.seal(_PAYLOAD)
    with pytest.raises(InvalidRequestState):
        codec.unseal("v2." + token.removeprefix(_TOKEN_PREFIX))


def test_a_kid_transplanted_onto_another_tokens_body_is_rejected() -> None:
    """Spec-mandated (basic/patterns/mrtr, server requirement 4): the kid is tag-bound; transplanting it fails."""
    raw_a = _decode_body(AESGCMRequestStateCodec([_KEY_A]).seal(_PAYLOAD))
    raw_b = _decode_body(AESGCMRequestStateCodec([_KEY_B]).seal(_PAYLOAD))
    assert raw_a[:_KID_LEN] != raw_b[:_KID_LEN]
    transplanted = _TOKEN_PREFIX + _b64u_nopad(raw_a[:_KID_LEN] + raw_b[_KID_LEN:])
    with pytest.raises(InvalidRequestState):
        AESGCMRequestStateCodec([_KEY_A, _KEY_B]).unseal(transplanted)


# -- RequestStateSecurity -----------------------------------------------------


def test_keys_and_codec_together_are_rejected_at_policy_construction() -> None:
    """SDK-defined: keys= and codec= are mutually exclusive."""
    with pytest.raises(ValueError) as exc:
        RequestStateSecurity(keys=[_KEY_A], codec=_StaticCodec())
    assert str(exc.value) == snapshot("RequestStateSecurity takes exactly one of keys= or codec=")


def test_a_policy_with_neither_keys_nor_codec_is_rejected() -> None:
    """SDK-defined: a policy must name its codec; an empty policy is a mistake, not a posture."""
    with pytest.raises(ValueError) as exc:
        RequestStateSecurity()
    assert str(exc.value) == snapshot("RequestStateSecurity takes exactly one of keys= or codec=")


@pytest.mark.parametrize(
    "ttl",
    [
        pytest.param(0.0, id="zero"),
        pytest.param(-600.0, id="negative"),
        pytest.param(float("nan"), id="nan"),
        pytest.param(float("inf"), id="inf"),
    ],
)
def test_a_non_positive_or_non_finite_ttl_is_rejected_at_policy_construction(ttl: float) -> None:
    """SDK-defined: zero, negative, NaN, and infinite ttl fail at construction for keys and ephemeral() alike."""
    with pytest.raises(ValueError, match="positive finite"):
        RequestStateSecurity(keys=[_KEY_A], ttl=ttl)
    with pytest.raises(ValueError, match="positive finite"):
        RequestStateSecurity.ephemeral(ttl=ttl)


def test_keys_produce_a_working_built_in_codec_on_the_policy() -> None:
    """SDK-defined: keys=[...] builds the built-in AES-GCM codec, exposed on .codec."""
    security = RequestStateSecurity(keys=[_KEY_A])
    assert isinstance(security.codec, AESGCMRequestStateCodec)
    assert security.codec.unseal(security.codec.seal(_PAYLOAD)) == _PAYLOAD


def test_a_custom_codec_is_stored_on_the_policy_as_is() -> None:
    """SDK-defined: codec=... stores the caller's object unwrapped."""
    codec = _StaticCodec()
    security = RequestStateSecurity(codec=codec)
    assert security.codec is codec
    assert codec.unseal(codec.seal(_PAYLOAD)) == _PAYLOAD


def test_ephemeral_policies_are_protected_and_mutually_unintelligible() -> None:
    """SDK-defined: ephemeral() protects under a process-local key, so a sibling instance rejects its tokens."""
    first = RequestStateSecurity.ephemeral()
    second = RequestStateSecurity.ephemeral()
    token = first.codec.seal(_PAYLOAD)
    assert first.codec.unseal(token) == _PAYLOAD
    with pytest.raises(InvalidRequestState):
        second.codec.unseal(token)


def test_the_policy_stores_an_explicit_audience_and_defaults_to_none() -> None:
    """SDK-defined: audience is stored as given; None defers to the server tier's `default_audience`."""
    assert RequestStateSecurity(keys=[_KEY_A]).audience is None
    assert RequestStateSecurity(keys=[_KEY_A], audience="svc").audience == "svc"
    assert RequestStateSecurity.ephemeral(audience="svc").audience == "svc"


def test_the_default_principal_binding_is_authenticated_principal() -> None:
    """SDK-defined: an unconfigured policy binds state to the authenticated OAuth client by default."""
    assert RequestStateSecurity(keys=[_KEY_A]).bind_principal is authenticated_principal


def test_an_explicit_principal_binding_callable_is_stored() -> None:
    """SDK-defined: a custom bind_principal callable is stored as given."""

    def tenant_binding(ctx: ServerRequestContext[Any, Any]) -> str | None:
        return "tenant-1"

    security = RequestStateSecurity(keys=[_KEY_A], bind_principal=tenant_binding)
    assert security.bind_principal is tenant_binding
    assert tenant_binding(_bare_context()) == "tenant-1"


# -- authenticated_principal ----------------------------------------------------


def test_authenticated_principal_is_none_without_an_auth_context() -> None:
    """SDK-defined: without an auth context the default binding derives no principal."""
    assert authenticated_principal(_bare_context()) is None


@pytest.mark.parametrize(
    ("token", "expected"),
    [
        pytest.param(
            AccessToken(token="at-1", client_id="client-123", scopes=[]),
            '["client-123",null,null]',
            id="client-only",
        ),
        pytest.param(
            AccessToken(token="at-2", client_id="client-123", scopes=[], subject="alice"),
            '["client-123",null,"alice"]',
            id="with-subject",
        ),
        pytest.param(
            AccessToken(
                token="at-3", client_id="client-123", scopes=[], subject="alice", claims={"iss": "https://as.example"}
            ),
            '["client-123","https://as.example","alice"]',
            id="with-issuer-and-subject",
        ),
    ],
)
def test_authenticated_principal_is_the_tokens_client_issuer_subject_identity(
    token: AccessToken, expected: str
) -> None:
    """SDK-defined: the default binding composes (client_id, issuer, subject), degrading per component."""
    reset = auth_context_var.set(AuthenticatedUser(token))
    try:
        assert authenticated_principal(_bare_context()) == expected
    finally:
        auth_context_var.reset(reset)


def test_authenticated_principal_distinguishes_two_subjects_of_one_client() -> None:
    """SDK-defined: two users of the same OAuth client are distinct principals when subjects are supplied."""
    alice = AccessToken(token="at-a", client_id="https://agent.example/client.json", scopes=[], subject="alice")
    bob = AccessToken(token="at-b", client_id="https://agent.example/client.json", scopes=[], subject="bob")
    principals: list[str | None] = []
    for token in (alice, bob):
        reset = auth_context_var.set(AuthenticatedUser(token))
        try:
            principals.append(authenticated_principal(_bare_context()))
        finally:
            auth_context_var.reset(reset)
    assert principals[0] != principals[1]


def test_authenticated_principal_uses_the_same_components_as_session_ownership() -> None:
    """SDK-defined: the binding and authorization_context derive from one principal_components source."""
    token = AccessToken(
        token="at-1", client_id="client-123", scopes=[], subject="alice", claims={"iss": "https://as.example"}
    )
    assert authorization_context(AuthenticatedUser(token)) == {
        "client_id": "client-123",
        "issuer": "https://as.example",
        "subject": "alice",
    }
    assert list(principal_components(token)) == ["client-123", "https://as.example", "alice"]