# Doc: https://docs.google.com/document/d/1jInrfvaD2DMsubehFJ-0k4KKwIe7xYhak6w6qqaq8jQ/edit#heading=h.7o12sd2rg8di # noqa
import json
import logging
import requests
from jwcrypto import jwt, jwk, jws

from django.contrib.auth.backends import ModelBackend
from django.contrib.auth import get_user_model
from django.db import transaction
from django.db.utils import IntegrityError

from university.utils.avatar import generate_lettere_avatar_url
from .models import SingleSignOnRecord
from .settings import sso_settings

logger = logging.getLogger(__name__)


class SsoAuthBackend(ModelBackend):
    """
    Authenticate user by pingcap cookie-sso flow.
    """

    def authenticate(self, request, sso_session_id):
        if not sso_session_id:
            return
        raw_user_data = self._get_sso_user_data(sso_session_id)
        if raw_user_data:
            user_data = json.loads(raw_user_data)
        else:
            return

        try:
            sso_record = SingleSignOnRecord.objects.get(external_id=user_data['sub'])
            sso_record.payload = raw_user_data
            sso_record.save(update_fields=['payload'])
            return sso_record.user if self.user_can_authenticate(sso_record.user) else None
        except SingleSignOnRecord.DoesNotExist:
            # Reslove concurrency issue when new user login
            # FIXME, ugly!!!
            try:
                with transaction.atomic():
                    new_user = get_user_model().objects.create(
                        username=user_data['username'],
                        email=user_data['email'] or None,
                        phone=user_data['phone'] or None,
                        avatar_url=generate_lettere_avatar_url(user_data['username'])
                    )
                    SingleSignOnRecord.objects.create(
                        user=new_user,
                        external_id=user_data['sub'],
                        payload=raw_user_data
                    )
                    return new_user
            except IntegrityError:
                import time
                time.sleep(0.2)
                return get_user_model().objects.get(username=user_data['username'])

    def _get_sso_user_data(self, sso_session_id):
        """
        Fetch user data from accounts.pingcap.com.
        check if signature is valid;
        check if jwt expired;
        """
        try:
            headers = {
                'Cookie': '{}={}'.format(sso_settings.SSO_COOKIE_NAME, sso_session_id),
                'Authorization': 'bearer {}'.format(sso_settings.IDP_API_TOKEN)
            }
            jwt_data = requests.get(
                sso_settings.IDP_API_ENDPOINT, headers=headers, timeout=1).json()['id_token']
            key = jwk.JWK.from_pem(sso_settings.JWT_PUBLIC_KEY.encode("utf8"))
            return jwt.JWT(key=key, jwt=jwt_data).claims
        except (jwt.JWTExpired, jws.InvalidJWSSignature, KeyError):
            pass
        except Exception:
            import traceback
            logger.error(traceback.format_exc())