import jwt as jwt_lib
import time


__all__ = ['Jwt', 'JwtDecodeError']


class JwtDecodeError(Exception):
    pass


class Jwt(object):
    """Base class for building a Json Web Token"""
    GENERATE = object()
    ALGORITHM = 'HS256'

    def __init__(self, secret_key, issuer, subject=None, algorithm=None, nbf=GENERATE,
                 ttl=3600, valid_until=None):
        self.secret_key = secret_key
        """:type str: The secret used to encode the JWT"""
        self.issuer = issuer
        """:type str: The issuer of this JWT"""
        self.subject = subject
        """:type str: The subject of this JWT, omitted from payload by default"""
        self.algorithm = algorithm or self.ALGORITHM
        """:type str: The algorithm used to encode the JWT, defaults to 'HS256'"""
        self.nbf = nbf
        """:type int: Time in secs since epoch before which this JWT is invalid. Defaults to now."""
        self.ttl = ttl
        """:type int: Time to live of the JWT in seconds, defaults to 1 hour"""
        self.valid_until = valid_until
        """:type int: Time in secs since epoch this JWT is valid for. Overrides ttl if provided."""

        self.__decoded_payload = None
        self.__decoded_headers = None

    def _generate_payload(self):
        """:rtype: dict the payload of the JWT to send"""
        raise NotImplementedError('Subclass must provide a payload.')

    def _generate_headers(self):
        """:rtype dict: Additional headers to include in the JWT, defaults to an empty dict"""
        return {}

    @classmethod
    def _from_jwt(cls, headers, payload, key=None):
        """
        Class specific implementation of from_jwt which should take jwt components and return
        and instance of this Class with jwt information loaded.
        :return: Jwt object containing the headers, payload and key
        """
        jwt = Jwt(
            secret_key=key,
            issuer=payload.get('iss', None),
            subject=payload.get('sub', None),
            algorithm=headers.get('alg', None),
            valid_until=payload.get('exp', None),
            nbf=payload.get('nbf', None),
        )
        jwt.__decoded_payload = payload
        jwt.__decoded_headers = headers
        return jwt

    @property
    def payload(self):
        if self.__decoded_payload:
            return self.__decoded_payload

        payload = self._generate_payload().copy()
        payload['iss'] = self.issuer
        payload['exp'] = int(time.time()) + self.ttl
        if self.nbf is not None:
            if self.nbf == self.GENERATE:
                payload['nbf'] = int(time.time())
            else:
                payload['nbf'] = self.nbf
        if self.valid_until:
            payload['exp'] = self.valid_until
        if self.subject:
            payload['sub'] = self.subject

        return payload

    @property
    def headers(self):
        if self.__decoded_headers:
            return self.__decoded_headers

        headers = self._generate_headers().copy()
        headers['typ'] = 'JWT'
        headers['alg'] = self.algorithm
        return headers

    def to_jwt(self, ttl=None):
        """
        Encode this JWT object into a JWT string
        :param int ttl: override the ttl configured in the constructor
        :rtype: str The JWT string
        """

        if not self.secret_key:
            raise ValueError('JWT does not have a signing key configured.')

        headers = self.headers.copy()

        payload = self.payload.copy()
        if ttl:
            payload['exp'] = int(time.time()) + ttl

        return jwt_lib.encode(payload, self.secret_key, algorithm=self.algorithm, headers=headers)

    @classmethod
    def from_jwt(cls, jwt, key=''):
        """
        Decode a JWT string into a Jwt object
        :param str jwt: JWT string
        :param Optional[str] key: key used to verify JWT signature, if not provided then validation
                                  is skipped.
        :raises JwtDecodeError if decoding JWT fails for any reason.
        :return: A DecodedJwt object containing the jwt information.
        """
        verify = True if key else False

        try:
            headers = jwt_lib.get_unverified_header(jwt)

            alg = headers.get('alg')
            if alg != cls.ALGORITHM:
                raise ValueError(f"Incorrect decoding algorithm {alg}, "
                                 f"expecting {cls.ALGORITHM}.")

            payload = jwt_lib.decode(jwt, key, algorithms=[cls.ALGORITHM], options={
                'verify_signature': verify,
                'verify_exp': True,
                'verify_nbf': True,
            })
        except Exception as e:
            raise JwtDecodeError(getattr(e, 'message', str(e)))

        return cls._from_jwt(headers, payload, key)

    def __str__(self):
        return '<JWT {}>'.format(self.to_jwt())