"""Custom JWT authentication that checks if the token has been blacklisted."""

from rest_framework_simplejwt.authentication import JWTAuthentication
from rest_framework_simplejwt.exceptions import InvalidToken
from rest_framework_simplejwt.token_blacklist.models import BlacklistedToken


class CustomJWTAuthentication(JWTAuthentication):
    """
    Extends JWTAuthentication to also check if the access token's
    associated refresh token has been blacklisted (i.e., user logged out).
    """

    def get_validated_token(self, raw_token):
        validated_token = super().get_validated_token(raw_token)

        # Check if this token's jti is blacklisted
        jti = validated_token.get("jti")
        user_id = validated_token.get("user_id")

        if jti:
            # Check if this specific token is blacklisted
            if BlacklistedToken.objects.filter(token__jti=jti).exists():
                raise InvalidToken("Token has been blacklisted")

        # Check if user has any blacklisted tokens after this token was issued
        # This effectively invalidates all access tokens when user logs out
        if user_id:
            token_iat = validated_token.get("iat")
            if token_iat:
                from datetime import datetime, timezone as dt_timezone

                issued_at = datetime.fromtimestamp(token_iat, tz=dt_timezone.utc)

                # If any refresh token for this user was blacklisted after this access token was issued,
                # consider this access token invalid
                blacklisted = BlacklistedToken.objects.filter(
                    token__user_id=user_id,
                    blacklisted_at__gte=issued_at,
                ).exists()

                if blacklisted:
                    raise InvalidToken("Token has been invalidated by logout")

        return validated_token