"""Tests for the identity-only gateway session token (mint/open, hostile-input totality)."""

from datetime import datetime, timedelta, timezone

import jwt
import pytest
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from pydantic import SecretStr, ValidationError

from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import (
    MAX_SESSION_TOKEN_BYTES,
    SESSION_ISSUER,
    SESSION_REFRESH_PREFIX,
    SESSION_REFRESH_TTL_SECONDS,
    SESSION_TOKEN_PREFIX,
    SESSION_TTL_SECONDS,
    AsymmetricSessionKeys,
    MintedSessionToken,
    NotASessionToken,
    OpenedSessionToken,
    SessionBadSignature,
    SessionExpired,
    SessionKeys,
    SessionMalformed,
    SessionPrincipal,
    SessionRotatedPublicKey,
    SessionTokenTooLarge,
    is_session_refresh_token,
    is_session_token,
    mint_session_refresh_token,
    mint_session_token,
    open_session_refresh_token,
    open_session_token,
    session_public_key_pem,
)


def _rsa_private_pem(bits: int = 2048) -> str:
    key = rsa.generate_private_key(public_exponent=65537, key_size=bits)
    return key.private_bytes(
        serialization.Encoding.PEM,
        serialization.PrivateFormat.PKCS8,
        serialization.NoEncryption(),
    ).decode()


_RSA_PEM_A = _rsa_private_pem()
_RSA_PEM_B = _rsa_private_pem()

NOW = datetime(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
KEYS = SessionKeys(signing_key=SecretStr("k" * 32))
OTHER_KEYS = SessionKeys(signing_key=SecretStr("x" * 32))
PRINCIPAL = SessionPrincipal(user_id="user-123", client_id="llm_client_abc")


def _mint_access() -> str:
    minted = mint_session_token(PRINCIPAL, KEYS, NOW)
    assert isinstance(minted, MintedSessionToken)
    return minted.token.get_secret_value()


def _mint_refresh() -> str:
    minted = mint_session_refresh_token(PRINCIPAL, KEYS, NOW)
    assert isinstance(minted, MintedSessionToken)
    return minted.token.get_secret_value()


def _sign_claims(payload: dict, prefix: str = SESSION_TOKEN_PREFIX, keys: SessionKeys = KEYS) -> str:
    return prefix + jwt.encode(payload, keys.signing_key.get_secret_value(), algorithm="HS256")


def _valid_claims(**overrides) -> dict:
    base = {
        "iss": SESSION_ISSUER,
        "iat": int(NOW.timestamp()),
        "exp": int((NOW + timedelta(seconds=600)).timestamp()),
        "jti": "jti-fixed",
        "kind": "session",
        "user_id": "user-123",
        "client_id": "llm_client_abc",
    }
    return {**base, **overrides}


def test_access_round_trip_recovers_principal_and_caps_ttl():
    minted = mint_session_token(PRINCIPAL, KEYS, NOW)
    assert isinstance(minted, MintedSessionToken)
    assert minted.expires_at == NOW + timedelta(seconds=SESSION_TTL_SECONDS)
    token = minted.token.get_secret_value()
    assert is_session_token(token)
    assert not is_session_refresh_token(token)
    opened = open_session_token(token, KEYS, NOW)
    assert isinstance(opened, OpenedSessionToken)
    assert opened.principal == PRINCIPAL


def test_refresh_round_trip_recovers_principal_and_caps_ttl():
    minted = mint_session_refresh_token(PRINCIPAL, KEYS, NOW)
    assert isinstance(minted, MintedSessionToken)
    assert minted.expires_at == NOW + timedelta(seconds=SESSION_REFRESH_TTL_SECONDS)
    token = minted.token.get_secret_value()
    assert is_session_refresh_token(token)
    opened = open_session_refresh_token(token, KEYS, NOW)
    assert isinstance(opened, OpenedSessionToken)
    assert opened.principal == PRINCIPAL


def test_access_token_reprefixed_as_refresh_is_rejected_by_signed_kind():
    body = _mint_access().removeprefix(SESSION_TOKEN_PREFIX)
    swapped = SESSION_REFRESH_PREFIX + body
    assert isinstance(open_session_refresh_token(swapped, KEYS, NOW), SessionMalformed)


def test_refresh_token_reprefixed_as_access_is_rejected_by_signed_kind():
    body = _mint_refresh().removeprefix(SESSION_REFRESH_PREFIX)
    swapped = SESSION_TOKEN_PREFIX + body
    assert isinstance(open_session_token(swapped, KEYS, NOW), SessionMalformed)


def test_refresh_token_is_not_an_access_token_at_the_edge():
    assert isinstance(open_session_token(_mint_refresh(), KEYS, NOW), NotASessionToken)


def test_expired_access_token_is_expired_not_malformed():
    token = _mint_access()
    at_expiry = NOW + timedelta(seconds=SESSION_TTL_SECONDS)
    assert isinstance(open_session_token(token, KEYS, at_expiry), SessionExpired)
    after = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1)
    assert isinstance(open_session_token(token, KEYS, after), SessionExpired)


def test_still_valid_one_second_before_expiry():
    token = _mint_access()
    just_before = NOW + timedelta(seconds=SESSION_TTL_SECONDS - 1)
    assert isinstance(open_session_token(token, KEYS, just_before), OpenedSessionToken)


def test_tampered_signature_is_bad_signature():
    token = _mint_access()
    tampered = token[:-2] + ("aa" if not token.endswith("aa") else "bb")
    assert isinstance(open_session_token(tampered, KEYS, NOW), SessionBadSignature)


def test_key_rotation_invalidates_outstanding_tokens():
    token = _mint_access()
    assert isinstance(open_session_token(token, OTHER_KEYS, NOW), SessionBadSignature)


@pytest.mark.parametrize(
    "candidate,expected",
    [
        ("sk-1234", NotASessionToken),
        ("llm_env_something", NotASessionToken),
        ("", NotASessionToken),
        (SESSION_TOKEN_PREFIX, SessionMalformed),
        (SESSION_TOKEN_PREFIX + "not-a-jwt", SessionMalformed),
        (SESSION_TOKEN_PREFIX + "\ud800garbage", SessionMalformed),
        (SESSION_TOKEN_PREFIX + "a" * (MAX_SESSION_TOKEN_BYTES + 1), SessionMalformed),
    ],
)
def test_hostile_candidates_never_raise(candidate, expected):
    assert isinstance(open_session_token(candidate, KEYS, NOW), expected)


def test_multibyte_candidate_over_byte_cap_but_under_char_cap_is_rejected():
    filler = "€" * (MAX_SESSION_TOKEN_BYTES // 3)
    candidate = SESSION_TOKEN_PREFIX + filler
    assert len(candidate) <= MAX_SESSION_TOKEN_BYTES
    assert isinstance(open_session_token(candidate, KEYS, NOW), SessionMalformed)


def test_alg_none_token_is_rejected():
    unsigned = jwt.api_jws.encode(b'{"iss":"litellm-mcp-gateway"}', key=None, algorithm="none")
    assert isinstance(open_session_token(SESSION_TOKEN_PREFIX + unsigned, KEYS, NOW), SessionMalformed)


@pytest.mark.parametrize(
    "claims",
    [
        _valid_claims(iss="wrong-issuer"),
        _valid_claims(exp=str(int((NOW + timedelta(seconds=600)).timestamp()))),
        _valid_claims(iat="evil"),
        _valid_claims(kind="access"),
        _valid_claims(user_id=""),
        _valid_claims(nbf=0),
        {k: v for k, v in _valid_claims().items() if k != "client_id"},
        {k: v for k, v in _valid_claims().items() if k != "exp"},
    ],
)
def test_signed_but_malformed_claims_are_rejected_without_raising(claims):
    token = _sign_claims(claims)
    assert isinstance(open_session_token(token, KEYS, NOW), SessionMalformed)


def test_signed_claims_with_exact_shape_open():
    token = _sign_claims(_valid_claims())
    opened = open_session_token(token, KEYS, NOW)
    assert isinstance(opened, OpenedSessionToken)
    assert opened.principal.user_id == "user-123"


def test_oversized_client_id_fails_mint_with_typed_error_not_truncation():
    principal = SessionPrincipal(user_id="user-123", client_id="c" * (MAX_SESSION_TOKEN_BYTES + 100))
    minted = mint_session_token(principal, KEYS, NOW)
    assert isinstance(minted, SessionTokenTooLarge)
    assert minted.max_bytes == MAX_SESSION_TOKEN_BYTES


def test_empty_principal_fields_rejected_at_construction():
    with pytest.raises(ValidationError):
        SessionPrincipal(user_id="", client_id="c")
    with pytest.raises(ValidationError):
        SessionPrincipal(user_id="u", client_id="")


def test_short_signing_key_rejected_at_construction():
    with pytest.raises(ValidationError):
        SessionKeys(signing_key=SecretStr("short"))


def test_two_mints_of_the_same_principal_are_distinct_tokens():
    first = mint_session_token(PRINCIPAL, KEYS, NOW)
    second = mint_session_token(PRINCIPAL, KEYS, NOW)
    assert isinstance(first, MintedSessionToken) and isinstance(second, MintedSessionToken)
    assert first.token.get_secret_value() != second.token.get_secret_value()


def test_minted_token_repr_never_leaks_value():
    minted = mint_session_token(PRINCIPAL, KEYS, NOW)
    assert isinstance(minted, MintedSessionToken)
    assert minted.token.get_secret_value() not in repr(minted)


def _decoded_claims(token: str, prefix: str) -> dict:
    return jwt.decode(
        token.removeprefix(prefix),
        KEYS.signing_key.get_secret_value(),
        algorithms=["HS256"],
        options={"verify_exp": False},
    )


def test_mcp_principal_wire_claims_carry_no_audience_or_team_keys():
    access_claims = _decoded_claims(_mint_access(), SESSION_TOKEN_PREFIX)
    refresh_claims = _decoded_claims(_mint_refresh(), SESSION_REFRESH_PREFIX)
    for claims in (access_claims, refresh_claims):
        assert "audience" not in claims
        assert "team_id" not in claims


def test_legacy_signed_claims_open_with_no_audience_and_no_team():
    opened = open_session_token(_sign_claims(_valid_claims()), KEYS, NOW)
    assert isinstance(opened, OpenedSessionToken)
    assert opened.principal.audience is None
    assert opened.principal.team_id is None


def test_proxy_api_audience_and_team_round_trip_through_the_refresh_token():
    principal = SessionPrincipal(user_id="user-123", client_id="llm_client_abc", audience="proxy_api", team_id="team-b")
    minted = mint_session_refresh_token(principal, KEYS, NOW)
    assert isinstance(minted, MintedSessionToken)
    token = minted.token.get_secret_value()
    claims = _decoded_claims(token, SESSION_REFRESH_PREFIX)
    assert claims["audience"] == "proxy_api"
    assert claims["team_id"] == "team-b"
    opened = open_session_refresh_token(token, KEYS, NOW)
    assert isinstance(opened, OpenedSessionToken)
    assert opened.principal == principal


def test_signed_claims_with_an_unknown_audience_are_rejected():
    token = _sign_claims(_valid_claims(audience="bogus"))
    assert isinstance(open_session_token(token, KEYS, NOW), SessionMalformed)


def test_signed_claims_with_a_non_string_team_are_rejected():
    token = _sign_claims(_valid_claims(team_id=42))
    assert isinstance(open_session_token(token, KEYS, NOW), SessionMalformed)


def test_principal_rejects_an_unknown_audience_at_construction():
    with pytest.raises(ValidationError):
        SessionPrincipal(user_id="user-123", client_id="llm_client_abc", audience="mcp")


RSA_KEYS = AsymmetricSessionKeys(private_key_pem=SecretStr(_RSA_PEM_A), kid="2026-01")
OTHER_RSA_KEYS = AsymmetricSessionKeys(private_key_pem=SecretStr(_RSA_PEM_B), kid="2025-06")


def test_rs256_access_round_trip_with_kid_and_alg_pinned_in_header():
    minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW)
    assert isinstance(minted, MintedSessionToken)
    token = minted.token.get_secret_value()
    header = jwt.get_unverified_header(token.removeprefix(SESSION_TOKEN_PREFIX))
    assert header["alg"] == "RS256"
    assert header["kid"] == "2026-01"
    opened = open_session_token(token, RSA_KEYS, NOW)
    assert isinstance(opened, OpenedSessionToken)
    assert opened.principal == PRINCIPAL


def test_rs256_refresh_round_trip():
    minted = mint_session_refresh_token(PRINCIPAL, RSA_KEYS, NOW)
    assert isinstance(minted, MintedSessionToken)
    token = minted.token.get_secret_value()
    opened = open_session_refresh_token(token, RSA_KEYS, NOW)
    assert isinstance(opened, OpenedSessionToken)
    assert opened.principal == PRINCIPAL


def test_rs256_token_verifies_with_public_key_only():
    minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW)
    assert isinstance(minted, MintedSessionToken)
    public_pem = session_public_key_pem(RSA_KEYS)
    assert "PUBLIC KEY" in public_pem
    assert "PRIVATE" not in public_pem
    claims = jwt.decode(
        minted.token.get_secret_value().removeprefix(SESSION_TOKEN_PREFIX),
        public_pem,
        algorithms=["RS256"],
        issuer=SESSION_ISSUER,
        options={"verify_exp": False},
    )
    assert claims["user_id"] == "user-123"


def test_rs256_tampered_signature_is_bad_signature():
    minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW)
    assert isinstance(minted, MintedSessionToken)
    token = minted.token.get_secret_value()
    tampered = token[:-2] + ("aa" if not token.endswith("aa") else "bb")
    assert isinstance(open_session_token(tampered, RSA_KEYS, NOW), SessionBadSignature)


def test_rs256_expired_token_is_expired():
    minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW)
    assert isinstance(minted, MintedSessionToken)
    after = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1)
    assert isinstance(open_session_token(minted.token.get_secret_value(), RSA_KEYS, after), SessionExpired)


def test_hs256_token_is_rejected_in_rs256_mode():
    assert isinstance(open_session_token(_mint_access(), RSA_KEYS, NOW), SessionBadSignature)


def test_hs256_token_claiming_the_current_kid_is_rejected_by_alg_pinning():
    token = SESSION_TOKEN_PREFIX + jwt.encode(
        _valid_claims(),
        KEYS.signing_key.get_secret_value(),
        algorithm="HS256",
        headers={"kid": RSA_KEYS.kid},
    )
    assert isinstance(open_session_token(token, RSA_KEYS, NOW), SessionMalformed)


def test_rs256_token_is_rejected_in_hs256_mode():
    minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW)
    assert isinstance(minted, MintedSessionToken)
    assert isinstance(open_session_token(minted.token.get_secret_value(), KEYS, NOW), SessionMalformed)


def test_rs256_token_from_an_unknown_kid_is_bad_signature():
    minted = mint_session_token(PRINCIPAL, OTHER_RSA_KEYS, NOW)
    assert isinstance(minted, MintedSessionToken)
    assert isinstance(open_session_token(minted.token.get_secret_value(), RSA_KEYS, NOW), SessionBadSignature)


def test_rs256_token_signed_by_a_foreign_key_claiming_the_current_kid_is_bad_signature():
    token = SESSION_TOKEN_PREFIX + jwt.encode(
        _valid_claims(),
        _RSA_PEM_B,
        algorithm="RS256",
        headers={"kid": RSA_KEYS.kid},
    )
    assert isinstance(open_session_token(token, RSA_KEYS, NOW), SessionBadSignature)


def test_alg_none_token_with_the_current_kid_is_rejected_in_rs256_mode():
    unsigned = jwt.api_jws.encode(
        b'{"iss":"litellm-mcp-gateway"}', key=None, algorithm="none", headers={"kid": RSA_KEYS.kid}
    )
    assert isinstance(open_session_token(SESSION_TOKEN_PREFIX + unsigned, RSA_KEYS, NOW), SessionMalformed)


def test_rotation_previous_public_key_still_verifies_until_removed():
    minted = mint_session_token(PRINCIPAL, OTHER_RSA_KEYS, NOW)
    assert isinstance(minted, MintedSessionToken)
    token = minted.token.get_secret_value()
    rotated = AsymmetricSessionKeys(
        private_key_pem=SecretStr(_RSA_PEM_A),
        kid="2026-01",
        previous_public_keys=(
            SessionRotatedPublicKey(kid="2025-06", public_key_pem=session_public_key_pem(OTHER_RSA_KEYS)),
        ),
    )
    opened = open_session_token(token, rotated, NOW)
    assert isinstance(opened, OpenedSessionToken)
    assert opened.principal == PRINCIPAL
    assert isinstance(open_session_token(token, RSA_KEYS, NOW), SessionBadSignature)


def test_rotation_window_still_enforces_expiry_and_tamper_on_the_previous_key():
    minted = mint_session_token(PRINCIPAL, OTHER_RSA_KEYS, NOW)
    assert isinstance(minted, MintedSessionToken)
    token = minted.token.get_secret_value()
    rotated = AsymmetricSessionKeys(
        private_key_pem=SecretStr(_RSA_PEM_A),
        kid="2026-01",
        previous_public_keys=(
            SessionRotatedPublicKey(kid="2025-06", public_key_pem=session_public_key_pem(OTHER_RSA_KEYS)),
        ),
    )
    after = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1)
    assert isinstance(open_session_token(token, rotated, after), SessionExpired)
    tampered = token[:-2] + ("aa" if not token.endswith("aa") else "bb")
    assert isinstance(open_session_token(tampered, rotated, NOW), SessionBadSignature)


def test_weak_or_garbage_private_key_pem_rejected_at_construction():
    with pytest.raises(ValidationError):
        AsymmetricSessionKeys(private_key_pem=SecretStr(_rsa_private_pem(bits=1024)), kid="weak")
    with pytest.raises(ValidationError):
        AsymmetricSessionKeys(private_key_pem=SecretStr("not a pem"), kid="junk")
    with pytest.raises(ValidationError):
        SessionRotatedPublicKey(kid="junk", public_key_pem="not a pem")
    with pytest.raises(ValidationError):
        SessionRotatedPublicKey(kid="private-half", public_key_pem=_RSA_PEM_A)


def test_weak_rotated_public_key_rejected_at_construction():
    weak_public = (
        serialization.load_pem_private_key(_rsa_private_pem(bits=1024).encode(), password=None)
        .public_key()
        .public_bytes(serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo)
        .decode()
    )
    with pytest.raises(ValidationError):
        SessionRotatedPublicKey(kid="2024-01", public_key_pem=weak_public)


def test_duplicate_kids_rejected_at_construction():
    previous = SessionRotatedPublicKey(kid="2025-06", public_key_pem=session_public_key_pem(OTHER_RSA_KEYS))
    with pytest.raises(ValidationError):
        AsymmetricSessionKeys(private_key_pem=SecretStr(_RSA_PEM_A), kid="2025-06", previous_public_keys=(previous,))
    with pytest.raises(ValidationError):
        AsymmetricSessionKeys(
            private_key_pem=SecretStr(_RSA_PEM_A), kid="2026-01", previous_public_keys=(previous, previous)
        )


def test_asymmetric_keys_repr_never_leaks_the_private_key():
    assert _RSA_PEM_A not in repr(RSA_KEYS)
