import binascii
import json
import zlib
from collections.abc import Mapping
from struct import pack

from . import jwk
from .backends import get_random_bytes
from .constants import ALGORITHMS, ZIPS
from .exceptions import JWEError, JWEParseError
from .utils import base64url_decode, base64url_encode, ensure_binary


def encrypt(plaintext, key, encryption=ALGORITHMS.A256GCM, algorithm=ALGORITHMS.DIR, zip=None, cty=None, kid=None):
    """Encrypts plaintext and returns a JWE cmpact serialization string.

    Args:
        plaintext (bytes): A bytes object to encrypt
        key (str or dict): The key(s) to use for encrypting the content. Can be
            individual JWK or JWK set.
        encryption (str, optional): The content encryption algorithm used to
            perform authenticated encryption on the plaintext to produce the
            ciphertext and the Authentication Tag.  Defaults to A256GCM.
        algorithm (str, optional): The cryptographic algorithm used
            to encrypt or determine the value of the CEK.  Defaults to dir.
        zip (str, optional): The compression algorithm) applied to the
            plaintext before encryption. Defaults to None.
        cty (str, optional): The media type for the secured content.
            See http://www.iana.org/assignments/media-types/media-types.xhtml
        kid (str, optional): Key ID for the provided key

    Returns:
        bytes: The string representation of the header, encrypted key,
            initialization vector, ciphertext, and authentication tag.

    Raises:
        JWEError: If there is an error signing the token.

    Examples:
        >>> from jose import jwe
        >>> jwe.encrypt('Hello, World!', 'asecret128bitkey', algorithm='dir', encryption='A128GCM')
        'eyJhbGciOiJkaXIiLCJlbmMiOiJBMTI4R0NNIn0..McILMB3dYsNJSuhcDzQshA.OfX9H_mcUpHDeRM4IA.CcnTWqaqxNsjT4eCaUABSg'

    """
    plaintext = ensure_binary(plaintext)  # Make sure it's bytes
    if algorithm not in ALGORITHMS.SUPPORTED:
        raise JWEError("Algorithm %s not supported." % algorithm)
    if encryption not in ALGORITHMS.SUPPORTED:
        raise JWEError("Algorithm %s not supported." % encryption)
    key = jwk.construct(key, algorithm)
    encoded_header = _encoded_header(algorithm, encryption, zip, cty, kid)

    plaintext = _compress(zip, plaintext)
    enc_cek, iv, cipher_text, auth_tag = _encrypt_and_auth(key, algorithm, encryption, zip, plaintext, encoded_header)

    jwe_string = _jwe_compact_serialize(encoded_header, enc_cek, iv, cipher_text, auth_tag)
    return jwe_string


def decrypt(jwe_str, key):
    """Decrypts a JWE compact serialized string and returns the plaintext.

    Args:
        jwe_str (str): A JWE to be decrypt.
        key (str or dict): A key to attempt to decrypt the payload with. Can be
            individual JWK or JWK set.

    Returns:
        bytes: The plaintext bytes, assuming the authentication tag is valid.

    Raises:
        JWEError: If there is an exception verifying the token.

    Examples:
        >>> from jose import jwe
        >>> jwe.decrypt(jwe_string, 'asecret128bitkey')
        'Hello, World!'
    """
    header, encoded_header, encrypted_key, iv, cipher_text, auth_tag = _jwe_compact_deserialize(jwe_str)

    # Verify that the implementation understands and can process all
    # fields that it is required to support, whether required by this
    # specification, by the algorithms being used, or by the "crit"
    # Header Parameter value, and that the values of those parameters
    # are also understood and supported.

    try:
        # Determine the Key Management Mode employed by the algorithm
        # specified by the "alg" (algorithm) Header Parameter.
        alg = header["alg"]
        enc = header["enc"]
        if alg not in ALGORITHMS.SUPPORTED:
            raise JWEError("Algorithm %s not supported." % alg)
        if enc not in ALGORITHMS.SUPPORTED:
            raise JWEError("Algorithm %s not supported." % enc)

    except KeyError:
        raise JWEParseError("alg and enc headers are required!")

    # Verify that the JWE uses a key known to the recipient.
    key = jwk.construct(key, alg)

    # When Direct Key Agreement or Key Agreement with Key Wrapping are
    # employed, use the key agreement algorithm to compute the value
    # of the agreed upon key.  When Direct Key Agreement is employed,
    # let the CEK be the agreed upon key.  When Key Agreement with Key
    # Wrapping is employed, the agreed upon key will be used to
    # decrypt the JWE Encrypted Key.
    #
    # When Key Wrapping, Key Encryption, or Key Agreement with Key
    # Wrapping are employed, decrypt the JWE Encrypted Key to produce
    # the CEK.  The CEK MUST have a length equal to that required for
    # the content encryption algorithm.  Note that when there are
    # multiple recipients, each recipient will only be able to decrypt
    # JWE Encrypted Key values that were encrypted to a key in that
    # recipient's possession.  It is therefore normal to only be able
    # to decrypt one of the per-recipient JWE Encrypted Key values to
    # obtain the CEK value.  Also, see Section 11.5 for security
    # considerations on mitigating timing attacks.
    if alg == ALGORITHMS.DIR:
        # When Direct Key Agreement or Direct Encryption are employed,
        # verify that the JWE Encrypted Key value is an empty octet
        # sequence.

        # Record whether the CEK could be successfully determined for this
        # recipient or not.
        cek_valid = encrypted_key == b""

        # When Direct Encryption is employed, let the CEK be the shared
        # symmetric key.
        cek_bytes = _get_key_bytes_from_key(key)
    else:
        try:
            cek_bytes = key.unwrap_key(encrypted_key)

            # Record whether the CEK could be successfully determined for this
            # recipient or not.
            cek_valid = True
        except NotImplementedError:
            raise JWEError(f"alg {alg} is not implemented")
        except Exception:
            # Record whether the CEK could be successfully determined for this
            # recipient or not.
            cek_valid = False

            # To mitigate the attacks described in RFC 3218 [RFC3218], the
            # recipient MUST NOT distinguish between format, padding, and length
            # errors of encrypted keys.  It is strongly recommended, in the event
            # of receiving an improperly formatted key, that the recipient
            # substitute a randomly generated CEK and proceed to the next step, to
            # mitigate timing attacks.
            cek_bytes = _get_random_cek_bytes_for_enc(enc)

    # Compute the Encoded Protected Header value BASE64URL(UTF8(JWE
    # Protected Header)).  If the JWE Protected Header is not present
    # (which can only happen when using the JWE JSON Serialization and
    # no "protected" member is present), let this value be the empty
    # string.
    protected_header = encoded_header

    # Let the Additional Authenticated Data encryption parameter be
    # ASCII(Encoded Protected Header).  However, if a JWE AAD value is
    # present (which can only be the case when using the JWE JSON
    # Serialization), instead let the Additional Authenticated Data
    # encryption parameter be ASCII(Encoded Protected Header || '.' ||
    # BASE64URL(JWE AAD)).
    aad = protected_header

    # Decrypt the JWE Ciphertext using the CEK, the JWE Initialization
    # Vector, the Additional Authenticated Data value, and the JWE
    # Authentication Tag (which is the Authentication Tag input to the
    # calculation) using the specified content encryption algorithm,
    # returning the decrypted plaintext and validating the JWE
    # Authentication Tag in the manner specified for the algorithm,
    # rejecting the input without emitting any decrypted output if the
    # JWE Authentication Tag is incorrect.
    try:
        plain_text = _decrypt_and_auth(cek_bytes, enc, cipher_text, iv, aad, auth_tag)
    except NotImplementedError:
        raise JWEError(f"enc {enc} is not implemented")
    except Exception as e:
        raise JWEError(e)

    # If a "zip" parameter was included, uncompress the decrypted
    # plaintext using the specified compression algorithm.
    if plain_text is not None:
        plain_text = _decompress(header.get("zip"), plain_text)

    return plain_text if cek_valid else None


def get_unverified_header(jwe_str):
    """Returns the decoded headers without verification of any kind.

    Args:
        jwe_str (str): A compact serialized JWE to decode the headers from.

    Returns:
        dict: The dict representation of the JWE headers.

    Raises:
        JWEError: If there is an exception decoding the JWE.
    """
    header = _jwe_compact_deserialize(jwe_str)[0]
    return header


def _decrypt_and_auth(cek_bytes, enc, cipher_text, iv, aad, auth_tag):
    """
    Decrypt and verify the data

    Args:
        cek_bytes (bytes): cek to derive encryption and possible auth key to
            verify the auth tag
        cipher_text (bytes): Encrypted data
        iv (bytes): Initialization vector (iv) used to encrypt data
        aad (bytes): Additional Authenticated Data used to verify the data
        auth_tag (bytes): Authentication ntag to verify the data

    Returns:
        (bytes): Decrypted data
    """
    # Decrypt the JWE Ciphertext using the CEK, the JWE Initialization
    # Vector, the Additional Authenticated Data value, and the JWE
    # Authentication Tag (which is the Authentication Tag input to the
    # calculation) using the specified content encryption algorithm,
    # returning the decrypted plaintext
    # and validating the JWE
    # Authentication Tag in the manner specified for the algorithm,
    if enc in ALGORITHMS.HMAC_AUTH_TAG:
        encryption_key, mac_key, key_len = _get_encryption_key_mac_key_and_key_length_from_cek(cek_bytes, enc)
        auth_tag_check = _auth_tag(cipher_text, iv, aad, mac_key, key_len)
    elif enc in ALGORITHMS.GCM:
        encryption_key = jwk.construct(cek_bytes, enc)
        auth_tag_check = auth_tag  # GCM check auth on decrypt
    else:
        raise NotImplementedError(f"enc {enc} is not implemented!")

    plaintext = encryption_key.decrypt(cipher_text, iv, aad, auth_tag)
    if auth_tag != auth_tag_check:
        raise JWEError("Invalid JWE Auth Tag")

    return plaintext


def _get_encryption_key_mac_key_and_key_length_from_cek(cek_bytes, enc):
    derived_key_len = len(cek_bytes) // 2
    mac_key_bytes = cek_bytes[0:derived_key_len]
    mac_key = _get_hmac_key(enc, mac_key_bytes)
    encryption_key_bytes = cek_bytes[-derived_key_len:]
    encryption_alg, _ = enc.split("-")
    encryption_key = jwk.construct(encryption_key_bytes, encryption_alg)
    return encryption_key, mac_key, derived_key_len


def _jwe_compact_deserialize(jwe_bytes):
    """
    Deserialize and verify the header and segments are appropriate.

    Args:
        jwe_bytes (bytes): The compact serialized JWE
    Returns:
        (dict, bytes, bytes, bytes, bytes, bytes)
    """

    # Base64url decode the encoded representations of the JWE
    # Protected Header, the JWE Encrypted Key, the JWE Initialization
    # Vector, the JWE Ciphertext, the JWE Authentication Tag, and the
    # JWE AAD, following the restriction that no line breaks,
    # whitespace, or other additional characters have been used.
    jwe_bytes = ensure_binary(jwe_bytes)
    try:
        header_segment, encrypted_key_segment, iv_segment, cipher_text_segment, auth_tag_segment = jwe_bytes.split(
            b".", 4
        )
        header_data = base64url_decode(header_segment)
    except ValueError:
        raise JWEParseError("Not enough segments")
    except (TypeError, binascii.Error):
        raise JWEParseError("Invalid header")

    # Verify that the octet sequence resulting from decoding the
    # encoded JWE Protected Header is a UTF-8-encoded representation
    # of a completely valid JSON object conforming to RFC 7159
    # [RFC7159]; let the JWE Protected Header be this JSON object.
    #
    # If using the JWE Compact Serialization, let the JOSE Header be
    # the JWE Protected Header.  Otherwise, when using the JWE JSON
    # Serialization, let the JOSE Header be the union of the members
    # of the JWE Protected Header, the JWE Shared Unprotected Header
    # and the corresponding JWE Per-Recipient Unprotected Header, all
    # of which must be completely valid JSON objects.  During this
    # step, verify that the resulting JOSE Header does not contain
    # duplicate Header Parameter names.  When using the JWE JSON
    # Serialization, this restriction includes that the same Header
    # Parameter name also MUST NOT occur in distinct JSON object
    # values that together comprise the JOSE Header.

    try:
        header = json.loads(header_data)
    except ValueError as e:
        raise JWEParseError(f"Invalid header string: {e}")

    if not isinstance(header, Mapping):
        raise JWEParseError("Invalid header string: must be a json object")

    try:
        encrypted_key = base64url_decode(encrypted_key_segment)
    except (TypeError, binascii.Error):
        raise JWEParseError("Invalid encrypted key")

    try:
        iv = base64url_decode(iv_segment)
    except (TypeError, binascii.Error):
        raise JWEParseError("Invalid IV")

    try:
        ciphertext = base64url_decode(cipher_text_segment)
    except (TypeError, binascii.Error):
        raise JWEParseError("Invalid cyphertext")

    try:
        auth_tag = base64url_decode(auth_tag_segment)
    except (TypeError, binascii.Error):
        raise JWEParseError("Invalid auth tag")

    return header, header_segment, encrypted_key, iv, ciphertext, auth_tag


def _encoded_header(alg, enc, zip, cty, kid):
    """
    Generate an appropriate JOSE header based on the values provided
    Args:
        alg (str): Key wrap/negotiation algorithm
        enc (str): Encryption algorithm
        zip (str): Compression method
        cty (str): Content type of the encrypted data
        kid (str): ID for the key used for the operation

    Returns:
        bytes: JSON object of header based on input
    """
    header = {"alg": alg, "enc": enc}
    if zip:
        header["zip"] = zip
    if cty:
        header["cty"] = cty
    if kid:
        header["kid"] = kid
    json_header = json.dumps(
        header,
        separators=(",", ":"),
        sort_keys=True,
    ).encode("utf-8")
    return base64url_encode(json_header)


def _big_endian(int_val):
    return pack("!Q", int_val)


def _encrypt_and_auth(key, alg, enc, zip, plaintext, aad):
    """
    Generate a content encryption key (cek) and initialization
    vector (iv) based on enc and alg, compress the plaintext based on zip,
    encrypt the compressed plaintext using the cek and iv based on enc

    Args:
        key (Key): The key provided for encryption
        alg (str): The algorithm use for key wrap/negotiation
        enc (str): The encryption algorithm with which to encrypt the plaintext
        zip (str): The compression algorithm with which to compress the plaintext
        plaintext (bytes): The data to encrypt
        aad (str): Additional authentication data utilized for generating an
                    auth tag

    Returns:
          (bytes, bytes, bytes, bytes): A tuple of the following data
                                 (key wrapped cek, iv, cipher text, auth tag)
    """
    try:
        cek_bytes, kw_cek = _get_cek(enc, alg, key)
    except NotImplementedError:
        raise JWEError(f"alg {alg} is not implemented")

    if enc in ALGORITHMS.HMAC_AUTH_TAG:
        encryption_key, mac_key, key_len = _get_encryption_key_mac_key_and_key_length_from_cek(cek_bytes, enc)
        iv, ciphertext, tag = encryption_key.encrypt(plaintext, aad)
        auth_tag = _auth_tag(ciphertext, iv, aad, mac_key, key_len)
    elif enc in ALGORITHMS.GCM:
        encryption_key = jwk.construct(cek_bytes, enc)
        iv, ciphertext, auth_tag = encryption_key.encrypt(plaintext, aad)
    else:
        raise NotImplementedError(f"enc {enc} is not implemented!")

    return kw_cek, iv, ciphertext, auth_tag


def _get_hmac_key(enc, mac_key_bytes):
    """
    Get an HMACKey for the provided encryption algorithm and key bytes

    Args:
        enc (str): Encryption algorithm
        mac_key_bytes (bytes): vytes for the HMAC key

    Returns:
         (HMACKey): The key to perform HMAC actions
    """
    _, hash_alg = enc.split("-")
    mac_key = jwk.construct(mac_key_bytes, hash_alg)
    return mac_key


def _compress(zip, plaintext):
    """
    Compress the plaintext based on the algorithm supplied

    Args:
        zip (str): Compression Algorithm
        plaintext (bytes): plaintext to compress

    Returns:
        (bytes): Compressed plaintext
    """
    if zip not in ZIPS.SUPPORTED:
        raise NotImplementedError("ZIP {} is not supported!")
    if zip is None:
        compressed = plaintext
    elif zip == ZIPS.DEF:
        compressed = zlib.compress(plaintext)
    else:
        raise NotImplementedError("ZIP {} is not implemented!")
    return compressed


def _decompress(zip, compressed):
    """
    Decompress the plaintext based on the algorithm supplied

    Args:
        zip (str): Compression Algorithm
        plaintext (bytes): plaintext to decompress

    Returns:
        (bytes): Compressed plaintext
    """
    if zip not in ZIPS.SUPPORTED:
        raise NotImplementedError("ZIP {} is not supported!")
    if zip is None:
        decompressed = compressed
    elif zip == ZIPS.DEF:
        decompressed = zlib.decompress(compressed)
    else:
        raise NotImplementedError("ZIP {} is not implemented!")
    return decompressed


def _get_cek(enc, alg, key):
    """
    Get the content encryption key

    Args:
        enc (str): Encryption algorithm
        alg (str): kwy wrap/negotiation algorithm
        key (Key): Key provided to encryption method

    Return:
        (bytes, bytes): Tuple of (cek bytes and wrapped cek)
    """
    if alg == ALGORITHMS.DIR:
        cek, wrapped_cek = _get_direct_key_wrap_cek(key)
    else:
        cek, wrapped_cek = _get_key_wrap_cek(enc, key)

    return cek, wrapped_cek


def _get_direct_key_wrap_cek(key):
    """
    Get the cek and wrapped cek from the encryption key direct

    Args:
        key (Key): Key provided to encryption method

    Return:
        (Key, bytes): Tuple of (cek Key object and wrapped cek)
    """
    # Get the JWK data to determine how to derive the cek
    jwk_data = key.to_dict()
    if jwk_data["kty"] == "oct":
        # Get the last half of an octal key as the cek
        cek_bytes = _get_key_bytes_from_key(key)
        wrapped_cek = b""
    else:
        raise NotImplementedError("JWK type {} not supported!".format(jwk_data["kty"]))
    return cek_bytes, wrapped_cek


def _get_key_bytes_from_key(key):
    """
    Get the raw key bytes from a Key object

    Args:
        key (Key): Key from which to extract the raw key bytes
    Returns:
        (bytes) key data
    """
    jwk_data = key.to_dict()
    encoded_key = jwk_data["k"]
    cek_bytes = base64url_decode(encoded_key)
    return cek_bytes


def _get_key_wrap_cek(enc, key):
    """_get_rsa_key_wrap_cek
    Get the content encryption key for RSA key wrap

    Args:
        enc (str): Encryption algorithm
        key (Key): Key provided to encryption method

    Returns:
        (Key, bytes): Tuple of (cek Key object and wrapped cek)
    """
    cek_bytes = _get_random_cek_bytes_for_enc(enc)
    wrapped_cek = key.wrap_key(cek_bytes)
    return cek_bytes, wrapped_cek


def _get_random_cek_bytes_for_enc(enc):
    """
    Get the random cek bytes based on the encryptionn algorithm

    Args:
        enc (str): Encryption algorithm

    Returns:
        (bytes) random bytes for cek key
    """
    if enc == ALGORITHMS.A128GCM:
        num_bits = 128
    elif enc == ALGORITHMS.A192GCM:
        num_bits = 192
    elif enc in (ALGORITHMS.A128CBC_HS256, ALGORITHMS.A256GCM):
        num_bits = 256
    elif enc == ALGORITHMS.A192CBC_HS384:
        num_bits = 384
    elif enc == ALGORITHMS.A256CBC_HS512:
        num_bits = 512
    else:
        raise NotImplementedError(f"{enc} not supported")
    cek_bytes = get_random_bytes(num_bits // 8)
    return cek_bytes


def _auth_tag(ciphertext, iv, aad, mac_key, tag_length):
    """
    Get ann auth tag from the provided data

    Args:
        ciphertext (bytes): Encrypted value
        iv (bytes): Initialization vector
        aad (bytes): Additional Authenticated Data
        mac_key (bytes): Key to use in generating the MAC
        tag_length (int): How log the tag should be

    Returns:
        (bytes) Auth tag
    """
    al = _big_endian(len(aad) * 8)
    auth_tag_input = aad + iv + ciphertext + al
    signature = mac_key.sign(auth_tag_input)
    auth_tag = signature[0:tag_length]
    return auth_tag


def _jwe_compact_serialize(encoded_header, encrypted_cek, iv, cipher_text, auth_tag):
    """
    Generate a compact serialized JWE

    Args:
        encoded_header (bytes): Base64 URL Encoded JWE header JSON
        encrypted_cek (bytes): Encrypted content encryption key (cek)
        iv (bytes): Initialization vector (IV)
        cipher_text (bytes): Cipher text
        auth_tag (bytes): JWE Auth Tag

    Returns:
        (str): JWE compact serialized string
    """
    cipher_text = ensure_binary(cipher_text)
    encoded_encrypted_cek = base64url_encode(encrypted_cek)
    encoded_iv = base64url_encode(iv)
    encoded_cipher_text = base64url_encode(cipher_text)
    encoded_auth_tag = base64url_encode(auth_tag)
    return (
        encoded_header
        + b"."
        + encoded_encrypted_cek
        + b"."
        + encoded_iv
        + b"."
        + encoded_cipher_text
        + b"."
        + encoded_auth_tag
    )