from rest_framework_simplejwt.authentication import JWTAuthentication
from .models import AdminTable
from utils.helper import failed_response
from rest_framework import status


class AdminJWTAuthentication(JWTAuthentication):
    def get_user(self, validated_token):
        admin_id = validated_token.get("user_id")
        if admin_id is None:
            raise failed_response(msg="User ID not found in token", code=status.HTTP_404_NOT_FOUND)
        try:
            admin = AdminTable.objects.get(id=admin_id)
        except AdminTable.DoesNotExist:
            raise failed_response(msg="User ID not found in token", code=status.HTTP_404_NOT_FOUND)
        return admin