"""Synthetic RS256 API-token example. Requires PyJWT[crypto].

This example uses a preapproved local key set. Production key publication,
bounded JWKS refresh, revocation and resource authorization are separate work.
No production tokens or private keys are included or printed.
"""

from datetime import UTC, datetime, timedelta

import jwt


def validate_access_token(token, approved_keys, *, issuer, audience):
    if len(token) > 16_384:
        raise ValueError("Token exceeds the example size limit")
    header = jwt.get_unverified_header(token)
    if header.get("alg") != "RS256" or header.get("typ") != "at+jwt":
        raise ValueError("Unexpected token profile")
    key_id = header.get("kid")
    if not isinstance(key_id, str) or key_id not in approved_keys:
        raise ValueError("Unknown verification key")
    claims = jwt.decode(
        token,
        approved_keys[key_id],
        algorithms=["RS256"],
        issuer=issuer,
        audience=audience,
        leeway=30,
        options={"require": ["exp", "iat", "iss", "aud", "sub", "client_id", "jti"]},
    )
    # This is an issuer-specific profile condition, not a generic JWT claim.
    if claims.get("token_use") != "access":
        raise ValueError("Unexpected token use")
    if not isinstance(claims.get("scope"), str):
        raise ValueError("Missing scope string")
    return claims


def self_test():
    from cryptography.hazmat.primitives.asymmetric import rsa

    private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    rotated_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    issuer = "https://issuer.example.test"
    audience = "https://orders.example.test"
    now = datetime.now(UTC)
    claims = {
        "iss": issuer,
        "aud": audience,
        "sub": "synthetic-user",
        "iat": now,
        "exp": now + timedelta(minutes=5),
        "token_use": "access",
        "scope": "orders.read",
        "client_id": "synthetic-client",
        "jti": "synthetic-token",
    }
    keys = {"current": private_key.public_key()}

    def encode(payload, *, key=private_key, kid="current", typ="at+jwt"):
        return jwt.encode(payload, key, algorithm="RS256", headers={"kid": kid, "typ": typ})

    def validate(token):
        return validate_access_token(token, keys, issuer=issuer, audience=audience)

    assert validate(encode(claims))["sub"] == "synthetic-user"
    cases = {
        "wrong issuer": encode({**claims, "iss": "https://other.example.test"}),
        "wrong audience": encode({**claims, "aud": "https://admin.example.test"}),
        "expired": encode({**claims, "exp": now - timedelta(minutes=1)}),
        "wrong use": encode({**claims, "token_use": "id"}),
        "wrong type": encode(claims, typ="JWT"),
        "unknown key": encode(claims, kid="unknown"),
        "untrusted signature": encode(claims, key=rotated_key),
        "missing expiry": encode({k: v for k, v in claims.items() if k != "exp"}),
        "missing client": encode({k: v for k, v in claims.items() if k != "client_id"}),
        "missing token identifier": encode({k: v for k, v in claims.items() if k != "jti"}),
    }
    for name, token in cases.items():
        try:
            validate(token)
        except (jwt.InvalidTokenError, ValueError):
            continue
        raise AssertionError(f"Accepted invalid case: {name}")
    keys["rotated"] = rotated_key.public_key()
    assert validate(encode(claims, key=rotated_key, kid="rotated"))["sub"] == "synthetic-user"
    print("Passed: valid token, 10 rejection cases, approved key rotation")


if __name__ == "__main__":
    self_test()