import requests
from jose import jwk, jwt
from jose.backends.base import Key
from jose.exceptions import JWTError


def get_kid(access_token: str) -> str:
    kid = jwt.get_unverified_header(access_token).get("kid")
    if kid is None:
        raise JWTError("'kid' not found in token header.")

    return kid


def get_key(kid: str, jwks_url: str) -> Key:
    jwks = requests.get(jwks_url, timeout=3).json()
    for jwk_data in jwks.get("keys", []):
        if jwk_data.get("kid") == kid:
            return jwk.construct(jwk_data)

    raise JWTError("Matching key not found.")
