import base64
import json
from typing import Optional

import jwt
import requests
from starlette import status

from src.error.exceptions import JWTTokenVerifyException
from src.system_var import environment


class JWTToken:
    id_token: str
    refresh_token: str
    verified_jwt: Optional[dict]

    def __init__(self, id_token: str, refresh_token: str, verified_jwt: dict = None) -> None:
        self.id_token = id_token
        self.refresh_token = refresh_token
        self.verified_jwt = verified_jwt

    @property
    def verified_token(self):
        if self.verified_jwt is None:
            raise JWTTokenVerifyException('検証されていないトークン')
        return self.verified_jwt

    @property
    def user_id(self):
        verified_token = self.verified_token

        user_id: str = None
        identities: dict = verified_token.get('identities')
        if identities is not None:
            # 一般ユーザーログインによる(SSO経由)
            user_id = identities[0]['userId']
            # <社員番号>@<ドメイン名>となっているため、@で分割
            user_id = user_id.split('@')[0]
        else:
            # メンテユーザーによるログイン
            user_id = verified_token.get('cognito:username')

        return user_id

    @classmethod
    def request(cls, code: str):
        """JWTをリクエストし、新たなインスタンスを返す

        Args:
            code (str): セキュリティコード

        Raises:
            JWTTokenVerifyException: 認証失敗

        Returns:
            JWTToken: JWTTokenのモデルクラス
        """
        token_url = f'{environment.COGNITO_AUTH_DOMAIN}/{environment.TOKEN_ENDPOINT}'
        request_params = {
            'grant_type': 'authorization_code',
            'client_id': environment.COGNITO_CLIENT_ID,
            'code': code,
            'redirect_uri': environment.COGNITO_REDIRECT_URI
        }

        message = bytes(f'{environment.COGNITO_CLIENT_ID}:{environment.COGNITO_CLIENT_SECRET}', 'utf8')
        auth_header_value = base64.b64encode(message).decode()
        request_headers = {
            'Content-Type': 'application/x-www-form-urlencoded',
            'Authorization': f'Basic {auth_header_value}'
        }

        res = requests.post(token_url, params=request_params, headers=request_headers)
        if res.status_code != status.HTTP_200_OK:
            raise JWTTokenVerifyException(res.text)

        token_response = json.loads(res.text)

        return cls(id_token=token_response['id_token'], refresh_token=token_response['refresh_token'])

    @classmethod
    def refresh(cls, refresh_token: str):
        """JWTをリフレッシュし、新たなインスタンスを返す

        Args:
            refresh_token (str): リフレッシュトークン

        Raises:
            JWTTokenVerifyException: 認証失敗

        Returns:
            JWTToken: JWTTokenのモデルクラス
        """
        token_url = f'{environment.COGNITO_AUTH_DOMAIN}/{environment.TOKEN_ENDPOINT}'
        request_params = {
            'grant_type': 'refresh_token',
            'client_id': environment.COGNITO_CLIENT_ID,
            'refresh_token': refresh_token,
            'redirect_uri': environment.COGNITO_REDIRECT_URI
        }

        message = bytes(f'{environment.COGNITO_CLIENT_ID}:{environment.COGNITO_CLIENT_SECRET}', 'utf8')
        auth_header_value = base64.b64encode(message).decode()
        request_headers = {
            'Content-Type': 'application/x-www-form-urlencoded',
            'Authorization': f'Basic {auth_header_value}'
        }

        res = requests.post(token_url, params=request_params, headers=request_headers)
        if res.status_code != status.HTTP_200_OK:
            raise JWTTokenVerifyException(res.text)

        token_response = json.loads(res.text)
        return cls(id_token=token_response['id_token'], refresh_token=refresh_token)

    def verify_token(self):
        if self.id_token is None:
            raise Exception('アクセストークンがない')

        issuer = f'https://cognito-idp.{environment.AWS_REGION}.amazonaws.com/{environment.COGNITO_USER_POOL_ID}'
        jwks_url = f'{issuer}/.well-known/jwks.json'

        jwks_client = jwt.PyJWKClient(jwks_url)
        signing_key = jwks_client.get_signing_key_from_jwt(self.id_token)

        try:
            verified_jwt = jwt.decode(
                self.id_token,
                signing_key.key,
                algorithms=['RS256'],
                audience=environment.COGNITO_CLIENT_ID,
                issuer=issuer,
                # Cognitoのサーバー時間とのズレにより、Issued atクレームの検証に失敗するパターンに対処する
                options={'verify_iat': False}
            )
        # 有効期限(exp)が切れた場合、トークンをリフレッシュする
        except jwt.ExpiredSignatureError:
            refreshed_jwt_token = JWTToken.refresh(self.refresh_token)
            return refreshed_jwt_token.verified_token()
        # 有効期限以外の検証に失敗した場合は例外とする
        except jwt.InvalidTokenError as e:
            raise JWTTokenVerifyException('Invalid token', e)
        # IDトークンを使用していることを検証する
        if verified_jwt['token_use'] != 'id':
            raise JWTTokenVerifyException('Invalid `token_use` claim, should be `id`.')

        return JWTToken(
            id_token=self.id_token,
            refresh_token=self.refresh_token,
            verified_jwt=verified_jwt
        )
