from __future__ import annotations

from typing import Any, Mapping

import jwt 
from jwt import PyJWTError

from ..config import CLAIM_AUD, CLAIM_CLIENT_ID, CLAIM_TOKEN_USE, CognitoAuthConfig
from ..exceptions import TokenError
from ..providers import build_jwt_client

class TokenVerifier:
    def __init__(self, config: CognitoAuthConfig, jwk_client: Any | None = None) -> None:
        self._config = config
        self._jwk_client = jwk_client or build_jwt_client(config)

    def verify(self, token: str) -> Mapping[str, Any]:
        try:
            signed_key = self._jwk_client.get_signing_key_from_jwt(token)
        except PyJWTError as e:
            raise TokenError(f'JWT verification failed: {e}', status_code=401) from e

        issuer = self._config.issuer()
        decode_kw: dict[str, Any] = {
            'algorithms': ['RS256'],
            'issuer': issuer,
            'options': {
                'require': ['exp', 'iss', 'token_use'],
                'verify_aud': False,
            }
        }

        try:
            payload = jwt.decode(token, signed_key.key, **decode_kw)
        except PyJWTError as e:
            raise TokenError(f'JWT verification failed: {e}', status_code=401) from e

        self._ensure_app_client_matches(payload)
        return payload

    def _ensure_app_client_matches(self, payload: Mapping[str, Any]) -> None:
        expected_client_id = self._config.app_client_id
        if not expected_client_id:
            return 

        token_use = payload.get(CLAIM_TOKEN_USE)
        if token_use == 'access':
            if payload.get(CLAIM_CLIENT_ID) != expected_client_id:
                raise TokenError(f'Access token client ID mismatch: {payload.get(CLAIM_CLIENT_ID)} != {expected_client_id}', status_code=401)
        elif token_use == 'id':
            if payload.get(CLAIM_AUD) != expected_client_id:
                raise TokenError(f'ID token client ID mismatch: {payload.get(CLAIM_AUD)} != {expected_client_id}', status_code=401)
        else:
            raise TokenError(f'Unknown token use: {token_use}', status_code=401)