from rest_framework_simplejwt.tokens import UntypedToken
from rest_framework_simplejwt.exceptions import TokenError
from rest_framework_simplejwt.authentication import JWTTokenUserAuthentication
from django.contrib.auth import get_user_model
from django.conf import settings

def get_user_from_validated_token(validated_token):
    user_id = validated_token.payload[settings.SIMPLE_JWT['USER_ID_CLAIM']]
    try:
        user = get_user_model().objects.get(id=user_id)
    except get_user_model().DoesNotExist:
        return None
    else:
        return user.profile, user

def validate_jwt_from_querystring(querystring):
    try:
        decoded_querystring = querystring.decode('utf-8')
        params = decoded_querystring.split('&')
        for param in params:
            key, value = param.split('=')
            if key == 'jwt':
                return validate_jwt(value)
    except:
        return False
    else:
        return False

def validate_jwt(token):
    try:
        return UntypedToken(token)
    except TokenError as e:
        return False

def validate(func):
    def _wrapper(self, content):
        validated_token = validate_jwt_from_querystring(self.scope['query_string'])
        if not validated_token:
            self.close()
        else:
            self.scope['profile'], self.scope['user'] = get_user_from_validated_token(validated_token)
            self.scope['user_id'] = self.scope['user'].id
            self.scope['profile_id'] = self.scope['profile'].id
            func(self, content)
    return _wrapper