"""JSON Web Key."""
import abc
import json
import logging
import math
from typing import (
    Any,
    Callable,
    Dict,
    Mapping,
    Optional,
    Sequence,
    Tuple,
    Type,
    Union,
)

import cryptography.exceptions
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec, rsa

import josepy.util
from josepy import errors, json_util, util

logger = logging.getLogger(__name__)


class JWK(json_util.TypedJSONObjectWithFields, metaclass=abc.ABCMeta):
    """JSON Web Key."""
    type_field_name = 'kty'
    TYPES: Dict[str, Type['JWK']] = {}
    cryptography_key_types: Tuple[Type[Any], ...] = ()
    """Subclasses should override."""

    required: Sequence[str] = NotImplemented
    """Required members of public key's representation as defined by JWK/JWA."""

    _thumbprint_json_dumps_params: Dict[str, Union[Optional[int], Sequence[str], bool]] = {
        # "no whitespace or line breaks before or after any syntactic
        # elements"
        'indent': None,
        'separators': (',', ':'),
        # "members ordered lexicographically by the Unicode [UNICODE]
        # code points of the member names"
        'sort_keys': True,
    }
    key: Any

    def thumbprint(self,
                   hash_function: Callable[[], hashes.HashAlgorithm] = hashes.SHA256) -> bytes:
        """Compute JWK Thumbprint.

        https://tools.ietf.org/html/rfc7638

        :returns: bytes

        """
        digest = hashes.Hash(hash_function(), backend=default_backend())
        digest.update(json.dumps(
            {k: v for k, v in self.to_json().items()
                if k in self.required},
            **self._thumbprint_json_dumps_params).encode())  # type: ignore[arg-type]
        return digest.finalize()

    @abc.abstractmethod
    def public_key(self) -> 'JWK':  # pragma: no cover
        """Generate JWK with public key.

        For symmetric cryptosystems, this would return ``self``.

        """
        raise NotImplementedError()

    @classmethod
    def _load_cryptography_key(cls, data: bytes, password: Optional[bytes] = None,
                               backend: Optional[Any] = None) -> Any:
        backend = default_backend() if backend is None else backend
        exceptions = {}

        # private key?
        loader_private: Any
        for loader_private in (serialization.load_pem_private_key,
                               serialization.load_der_private_key):
            try:
                return loader_private(data, password, backend)
            except (ValueError, TypeError,
                    cryptography.exceptions.UnsupportedAlgorithm) as error:
                exceptions[str(loader_private)] = error

        # public key?
        loader_public: Any
        for loader_public in (serialization.load_pem_public_key,
                              serialization.load_der_public_key):
            try:
                return loader_public(data, backend)
            except (ValueError,
                    cryptography.exceptions.UnsupportedAlgorithm) as error:
                exceptions[str(loader_public)] = error

        # no luck
        raise errors.Error('Unable to deserialize key: {0}'.format(exceptions))

    @classmethod
    def load(cls, data: bytes, password: Optional[bytes] = None,
             backend: Optional[Any] = None) -> 'JWK':
        """Load serialized key as JWK.

        :param str data: Public or private key serialized as PEM or DER.
        :param str password: Optional password.
        :param backend: A `.PEMSerializationBackend` and
            `.DERSerializationBackend` provider.

        :raises errors.Error: if unable to deserialize, or unsupported
            JWK algorithm

        :returns: JWK of an appropriate type.
        :rtype: `JWK`

        """
        try:
            key = cls._load_cryptography_key(data, password, backend)
        except errors.Error as error:
            logger.debug('Loading symmetric key, asymmetric failed: %s', error)
            return JWKOct(key=data)

        if cls.typ is not NotImplemented and not isinstance(key, cls.cryptography_key_types):
            raise errors.Error('Unable to deserialize {0} into {1}'.format(
                key.__class__, cls.__class__))
        for jwk_cls in cls.TYPES.values():
            if isinstance(key, jwk_cls.cryptography_key_types):
                return jwk_cls(key=key)
        raise errors.Error('Unsupported algorithm: {0}'.format(key.__class__))


@JWK.register
class JWKOct(JWK):
    """Symmetric JWK."""
    typ = 'oct'
    __slots__ = ('key',)
    required = ('k', JWK.type_field_name)
    key: bytes

    def fields_to_partial_json(self) -> Dict[str, str]:
        # TODO: An "alg" member SHOULD also be present to identify the
        # algorithm intended to be used with the key, unless the
        # application uses another means or convention to determine
        # the algorithm used.
        return {'k': json_util.encode_b64jose(self.key)}

    @classmethod
    def fields_from_json(cls, jobj: Mapping[str, Any]) -> 'JWKOct':
        return cls(key=json_util.decode_b64jose(jobj['k']))

    def public_key(self) -> 'JWKOct':
        return self


@JWK.register
class JWKRSA(JWK):
    """RSA JWK.

    :ivar key: :class:`~cryptography.hazmat.primitives.asymmetric.rsa.RSAPrivateKey`
        or :class:`~cryptography.hazmat.primitives.asymmetric.rsa.RSAPublicKey` wrapped
        in :class:`~josepy.util.ComparableRSAKey`

    """
    typ = 'RSA'
    cryptography_key_types = (rsa.RSAPublicKey, rsa.RSAPrivateKey)
    __slots__ = ('key',)
    required = ('e', JWK.type_field_name, 'n')
    key: josepy.util.ComparableRSAKey

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        if 'key' in kwargs and not isinstance(
                kwargs['key'], util.ComparableRSAKey):
            kwargs['key'] = util.ComparableRSAKey(kwargs['key'])
        super().__init__(*args, **kwargs)

    @classmethod
    def _encode_param(cls, data: int) -> str:
        """Encode Base64urlUInt.
        :type data: long
        :rtype: unicode
        """
        length = max(data.bit_length(), 8)  # decoding 0
        length = math.ceil(length / 8)
        return json_util.encode_b64jose(data.to_bytes(byteorder="big", length=length))

    @classmethod
    def _decode_param(cls, data: str) -> int:
        """Decode Base64urlUInt."""
        try:
            binary = json_util.decode_b64jose(data)
            if not binary:
                raise errors.DeserializationError()
            return int.from_bytes(binary, byteorder="big")
        except ValueError:  # invalid literal for long() with base 16
            raise errors.DeserializationError()

    def public_key(self) -> 'JWKRSA':
        return type(self)(key=self.key.public_key())

    @classmethod
    def fields_from_json(cls, jobj: Mapping[str, Any]) -> 'JWKRSA':
        n, e = (cls._decode_param(jobj[x]) for x in ('n', 'e'))
        public_numbers = rsa.RSAPublicNumbers(e=e, n=n)

        # public key
        if 'd' not in jobj:
            return cls(key=public_numbers.public_key(default_backend()))

        # private key
        d = cls._decode_param(jobj['d'])
        if ('p' in jobj or 'q' in jobj or 'dp' in jobj or
                'dq' in jobj or 'qi' in jobj or 'oth' in jobj):
            # "If the producer includes any of the other private
            # key parameters, then all of the others MUST be
            # present, with the exception of "oth", which MUST
            # only be present when more than two prime factors
            # were used."
            p, q, dp, dq, qi, = all_params = tuple(
                jobj.get(x) for x in ('p', 'q', 'dp', 'dq', 'qi'))
            if tuple(param for param in all_params if param is None):
                raise errors.Error(
                    'Some private parameters are missing: {0}'.format(
                        all_params))
            p, q, dp, dq, qi = tuple(
                cls._decode_param(str(x)) for x in all_params)

            # TODO: check for oth
        else:
            # cryptography>=0.8
            p, q = rsa.rsa_recover_prime_factors(n, e, d)
            dp = rsa.rsa_crt_dmp1(d, p)
            dq = rsa.rsa_crt_dmq1(d, q)
            qi = rsa.rsa_crt_iqmp(p, q)

        key = rsa.RSAPrivateNumbers(
            p, q, d, dp, dq, qi, public_numbers).private_key(
                default_backend())

        return cls(key=key)

    def fields_to_partial_json(self) -> Dict[str, Any]:
        if isinstance(self.key._wrapped, rsa.RSAPublicKey):
            numbers = self.key.public_numbers()
            params = {
                'n': numbers.n,
                'e': numbers.e,
            }
        else:  # rsa.RSAPrivateKey
            private = self.key.private_numbers()
            public = self.key.public_key().public_numbers()
            params = {
                'n': public.n,
                'e': public.e,
                'd': private.d,
                'p': private.p,
                'q': private.q,
                'dp': private.dmp1,
                'dq': private.dmq1,
                'qi': private.iqmp,
            }
        return {key: self._encode_param(value)
                for key, value in params.items()}


@JWK.register
class JWKEC(JWK):
    """EC JWK.

    :ivar key: :class:`~cryptography.hazmat.primitives.asymmetric.ec.EllipticCurvePrivateKey`
        or :class:`~cryptography.hazmat.primitives.asymmetric.ec.EllipticCurvePublicKey` wrapped
        in :class:`~josepy.util.ComparableECKey`

    """
    typ = 'EC'
    __slots__ = ('key',)
    cryptography_key_types = (
        ec.EllipticCurvePublicKey, ec.EllipticCurvePrivateKey)
    required = ('crv', JWK.type_field_name, 'x', 'y')
    key: josepy.util.ComparableECKey

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        if 'key' in kwargs and not isinstance(
                kwargs['key'], util.ComparableECKey):
            kwargs['key'] = util.ComparableECKey(kwargs['key'])
        super().__init__(*args, **kwargs)

    @classmethod
    def _encode_param(cls, data: int, length: int) -> str:
        """Encode Base64urlUInt.
        :type data: long
        :type key_size: long
        :rtype: unicode
        """
        return json_util.encode_b64jose(data.to_bytes(byteorder="big", length=length))

    @classmethod
    def _decode_param(cls, data: str, name: str, valid_length: int) -> int:
        """Decode Base64urlUInt."""
        try:
            binary = json_util.decode_b64jose(data)
            if len(binary) != valid_length:
                raise errors.DeserializationError(
                    f'Expected parameter "{name}" to be {valid_length} bytes '
                    f'after base64-decoding; got {len(binary)} bytes instead'
                )
            return int.from_bytes(binary, byteorder="big")
        except ValueError:  # invalid literal for long() with base 16
            raise errors.DeserializationError()

    @classmethod
    def _curve_name_to_crv(cls, curve_name: str) -> str:
        if curve_name == 'secp256r1':
            return 'P-256'
        if curve_name == 'secp384r1':
            return 'P-384'
        if curve_name == 'secp521r1':
            return 'P-521'
        raise errors.SerializationError()

    @classmethod
    def _crv_to_curve(cls, crv: str) -> ec.EllipticCurve:
        # crv is case-sensitive
        if crv == 'P-256':
            return ec.SECP256R1()
        if crv == 'P-384':
            return ec.SECP384R1()
        if crv == 'P-521':
            return ec.SECP521R1()
        raise errors.DeserializationError()

    @classmethod
    def expected_length_for_curve(cls, curve: ec.EllipticCurve) -> int:
        if isinstance(curve, ec.SECP256R1):
            return 32
        elif isinstance(curve, ec.SECP384R1):
            return 48
        elif isinstance(curve, ec.SECP521R1):
            return 66
        raise ValueError(f'Unexpected curve: {curve}')

    def fields_to_partial_json(self) -> Dict[str, Any]:
        params = {}
        if isinstance(self.key._wrapped, ec.EllipticCurvePublicKey):
            public = self.key.public_numbers()
        elif isinstance(self.key._wrapped, ec.EllipticCurvePrivateKey):
            private = self.key.private_numbers()
            public = self.key.public_key().public_numbers()
            params['d'] = private.private_value
        else:
            raise errors.SerializationError(
                'Supplied key is neither of type EllipticCurvePublicKey '
                'nor EllipticCurvePrivateKey'
            )
        params['x'] = public.x
        params['y'] = public.y
        params = {
            key: self._encode_param(value, self.expected_length_for_curve(public.curve))
            for key, value in params.items()
        }
        params['crv'] = self._curve_name_to_crv(public.curve.name)
        return params

    @classmethod
    def fields_from_json(cls, jobj: Mapping[str, Any]) -> 'JWKEC':
        curve = cls._crv_to_curve(jobj['crv'])
        expected_length = cls.expected_length_for_curve(curve)
        x, y = (cls._decode_param(jobj[n], n, expected_length) for n in ('x', 'y'))
        public_numbers = ec.EllipticCurvePublicNumbers(x=x, y=y, curve=curve)

        # private key
        if 'd' not in jobj:
            return cls(key=public_numbers.public_key(default_backend()))

        # private key
        d = cls._decode_param(jobj['d'], 'd', expected_length)
        key = ec.EllipticCurvePrivateNumbers(d, public_numbers).private_key(default_backend())
        return cls(key=key)

    def public_key(self) -> 'JWKEC':
        # Unlike RSAPrivateKey, EllipticCurvePrivateKey does not contain public_key()
        if hasattr(self.key, 'public_key'):
            key = self.key.public_key()
        else:
            key = self.key.public_numbers().public_key(default_backend())
        return type(self)(key=key)