import jwt
from rest_framework import exceptions
from django.utils.encoding import smart_str
from users.models import UserSession,User
from organisations.models import Organisation
from django.utils.translation import gettext as _
from questbanker_api.utils import set_logged_in_user_key
from rest_framework.authentication import get_authorization_header
from rest_framework_jwt.authentication import JSONWebTokenAuthentication
from rest_framework_jwt.settings import api_settings
import json

jwt_decode_handler = api_settings.JWT_DECODE_HANDLER
jwt_get_username_from_payload = api_settings.JWT_PAYLOAD_GET_USERNAME_HANDLER

class CustomJSONWebTokenAuthentication(JSONWebTokenAuthentication):
    """
    Token based authentication using the JSON Web Token standard.
    """
    def authenticate(self, request):
        jwt_value = self.get_jwt_value(request)
        if jwt_value is None:
            return None

        try:
            payload = jwt_decode_handler(jwt_value)
        except jwt.ExpiredSignature:
            msg = _('Signature has expired.')
            raise exceptions.AuthenticationFailed(msg)
        except jwt.DecodeError:
            msg = _('Error decoding signature.')
            raise exceptions.AuthenticationFailed(msg)
        except jwt.InvalidTokenError:
            raise exceptions.AuthenticationFailed()

        user = self.authenticate_credentials(payload)

        # Custom session/working hours checks
        try:
            token_expired = True
            matched_session = None
            user_sessions = UserSession.objects.filter(user=user)
            if user_sessions:
                for user_session in user_sessions:
                    if user_session.session_token:
                        if user_session.session_token in jwt_value.decode('utf-8') if isinstance(jwt_value, bytes) else jwt_value:
                            matched_session = user_session
                            current_url = request.build_absolute_uri()
                            if 'verify-otp' in str(current_url) or 'verify-customer-otp' in str(current_url):
                                token_expired = False
                            elif user_session.allow_access == True:
                                token_expired = False
                            set_logged_in_user_key(user_session.session_token)
                            break

            if not matched_session:
                raise exceptions.AuthenticationFailed('Signature has expired.')

            organisation_id = None
            session_data = matched_session.data if isinstance(matched_session.data, dict) else {}
            if session_data:
                organisation_id = session_data.get('organisation_id')
            if not organisation_id and getattr(user, 'user_organisation_branch_id', None):
                organisation_id = getattr(user.user_organisation_branch, 'branch_organisation_id', None)

            organisation = Organisation.objects.filter(pk=organisation_id).first() if organisation_id else None
            from license.helpers import get_organisation_access_state
            access_state = get_organisation_access_state(organisation=organisation)
            if not access_state['allowed']:
                raise exceptions.AuthenticationFailed(access_state['message'])

            if token_expired:
                raise exceptions.AuthenticationFailed('Signature has expired.')
        except exceptions.AuthenticationFailed:
            raise
        except Exception:
            raise exceptions.AuthenticationFailed('Authentication failed.')

        return (user, jwt_value)

    def get_jwt_value(self, request):
        auth = get_authorization_header(request).split()
        valid_prefixes = {'bearer', 'jwt', 'token'}

        if auth:
            if len(auth) == 1:
                # Support raw token in Authorization header (no prefix).
                return auth[0]
            elif len(auth) == 2:
                prefix = smart_str(auth[0]).lower()
                if prefix in valid_prefixes:
                    return auth[1]
                return None
            else:
                msg = _('Invalid Authorization header. Credentials string '
                        'should not contain spaces.')
                raise exceptions.AuthenticationFailed(msg)

        # Fallback for clients that send token in X-Access-Token.
        x_access_token = request.META.get('HTTP_X_ACCESS_TOKEN', '')
        if x_access_token:
            return x_access_token

        return None