import os
import struct
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.keywrap import (
    aes_key_wrap,
    aes_key_unwrap
)
from cryptography.hazmat.primitives.ciphers import Cipher
from cryptography.hazmat.primitives.ciphers.algorithms import AES
from cryptography.hazmat.primitives.ciphers.modes import GCM
from cryptography.hazmat.primitives.kdf.concatkdf import ConcatKDFHash
from authlib.common.encoding import (
    to_bytes, to_native,
    urlsafe_b64decode,
    urlsafe_b64encode
)
from authlib.jose.rfc7516 import JWEAlgorithm
from .rsa_key import RSAKey
from .ec_key import ECKey
from .oct_key import OctKey


class DirectAlgorithm(JWEAlgorithm):
    name = 'dir'
    description = 'Direct use of a shared symmetric key'

    def prepare_key(self, raw_data):
        return OctKey.import_key(raw_data)

    def generate_preset(self, enc_alg, key):
        return {}

    def wrap(self, enc_alg, headers, key, preset=None):
        cek = key.get_op_key('encrypt')
        if len(cek) * 8 != enc_alg.CEK_SIZE:
            raise ValueError('Invalid "cek" length')
        return {'ek': b'', 'cek': cek}

    def unwrap(self, enc_alg, ek, headers, key):
        cek = key.get_op_key('decrypt')
        if len(cek) * 8 != enc_alg.CEK_SIZE:
            raise ValueError('Invalid "cek" length')
        return cek


class RSAAlgorithm(JWEAlgorithm):
    #: A key of size 2048 bits or larger MUST be used with these algorithms
    #: RSA1_5, RSA-OAEP, RSA-OAEP-256
    key_size = 2048

    def __init__(self, name, description, pad_fn):
        self.name = name
        self.description = description
        self.padding = pad_fn

    def prepare_key(self, raw_data):
        return RSAKey.import_key(raw_data)

    def generate_preset(self, enc_alg, key):
        cek = enc_alg.generate_cek()
        return {'cek': cek}

    def wrap(self, enc_alg, headers, key, preset=None):
        if preset and 'cek' in preset:
            cek = preset['cek']
        else:
            cek = enc_alg.generate_cek()

        op_key = key.get_op_key('wrapKey')
        if op_key.key_size < self.key_size:
            raise ValueError('A key of size 2048 bits or larger MUST be used')
        ek = op_key.encrypt(cek, self.padding)
        return {'ek': ek, 'cek': cek}

    def unwrap(self, enc_alg, ek, headers, key):
        # it will raise ValueError if failed
        op_key = key.get_op_key('unwrapKey')
        cek = op_key.decrypt(ek, self.padding)
        if len(cek) * 8 != enc_alg.CEK_SIZE:
            raise ValueError('Invalid "cek" length')
        return cek


class AESAlgorithm(JWEAlgorithm):
    def __init__(self, key_size):
        self.name = f'A{key_size}KW'
        self.description = f'AES Key Wrap using {key_size}-bit key'
        self.key_size = key_size

    def prepare_key(self, raw_data):
        return OctKey.import_key(raw_data)

    def generate_preset(self, enc_alg, key):
        cek = enc_alg.generate_cek()
        return {'cek': cek}

    def _check_key(self, key):
        if len(key) * 8 != self.key_size:
            raise ValueError(
                f'A key of size {self.key_size} bits is required.')

    def wrap_cek(self, cek, key):
        op_key = key.get_op_key('wrapKey')
        self._check_key(op_key)
        ek = aes_key_wrap(op_key, cek, default_backend())
        return {'ek': ek, 'cek': cek}

    def wrap(self, enc_alg, headers, key, preset=None):
        if preset and 'cek' in preset:
            cek = preset['cek']
        else:
            cek = enc_alg.generate_cek()
        return self.wrap_cek(cek, key)

    def unwrap(self, enc_alg, ek, headers, key):
        op_key = key.get_op_key('unwrapKey')
        self._check_key(op_key)
        cek = aes_key_unwrap(op_key, ek, default_backend())
        if len(cek) * 8 != enc_alg.CEK_SIZE:
            raise ValueError('Invalid "cek" length')
        return cek


class AESGCMAlgorithm(JWEAlgorithm):
    EXTRA_HEADERS = frozenset(['iv', 'tag'])

    def __init__(self, key_size):
        self.name = f'A{key_size}GCMKW'
        self.description = f'Key wrapping with AES GCM using {key_size}-bit key'
        self.key_size = key_size

    def prepare_key(self, raw_data):
        return OctKey.import_key(raw_data)

    def generate_preset(self, enc_alg, key):
        cek = enc_alg.generate_cek()
        return {'cek': cek}

    def _check_key(self, key):
        if len(key) * 8 != self.key_size:
            raise ValueError(
                f'A key of size {self.key_size} bits is required.')

    def wrap(self, enc_alg, headers, key, preset=None):
        if preset and 'cek' in preset:
            cek = preset['cek']
        else:
            cek = enc_alg.generate_cek()

        op_key = key.get_op_key('wrapKey')
        self._check_key(op_key)

        #: https://tools.ietf.org/html/rfc7518#section-4.7.1.1
        #: The "iv" (initialization vector) Header Parameter value is the
        #: base64url-encoded representation of the 96-bit IV value
        iv_size = 96
        iv = os.urandom(iv_size // 8)

        cipher = Cipher(AES(op_key), GCM(iv), backend=default_backend())
        enc = cipher.encryptor()
        ek = enc.update(cek) + enc.finalize()

        h = {
            'iv': to_native(urlsafe_b64encode(iv)),
            'tag': to_native(urlsafe_b64encode(enc.tag))
        }
        return {'ek': ek, 'cek': cek, 'header': h}

    def unwrap(self, enc_alg, ek, headers, key):
        op_key = key.get_op_key('unwrapKey')
        self._check_key(op_key)

        iv = headers.get('iv')
        if not iv:
            raise ValueError('Missing "iv" in headers')

        tag = headers.get('tag')
        if not tag:
            raise ValueError('Missing "tag" in headers')

        iv = urlsafe_b64decode(to_bytes(iv))
        tag = urlsafe_b64decode(to_bytes(tag))

        cipher = Cipher(AES(op_key), GCM(iv, tag), backend=default_backend())
        d = cipher.decryptor()
        cek = d.update(ek) + d.finalize()
        if len(cek) * 8 != enc_alg.CEK_SIZE:
            raise ValueError('Invalid "cek" length')
        return cek


class ECDHESAlgorithm(JWEAlgorithm):
    EXTRA_HEADERS = ['epk', 'apu', 'apv']
    ALLOWED_KEY_CLS = ECKey

    # https://tools.ietf.org/html/rfc7518#section-4.6
    def __init__(self, key_size=None):
        if key_size is None:
            self.name = 'ECDH-ES'
            self.description = 'ECDH-ES in the Direct Key Agreement mode'
        else:
            self.name = f'ECDH-ES+A{key_size}KW'
            self.description = (
                'ECDH-ES using Concat KDF and CEK wrapped '
                'with A{}KW').format(key_size)
        self.key_size = key_size
        self.aeskw = AESAlgorithm(key_size)

    def prepare_key(self, raw_data):
        if isinstance(raw_data, self.ALLOWED_KEY_CLS):
            return raw_data
        return ECKey.import_key(raw_data)

    def generate_preset(self, enc_alg, key):
        epk = self._generate_ephemeral_key(key)
        h = self._prepare_headers(epk)
        preset = {'epk': epk, 'header': h}
        if self.key_size is not None:
            cek = enc_alg.generate_cek()
            preset['cek'] = cek
        return preset

    def compute_fixed_info(self, headers, bit_size):
        # AlgorithmID
        if self.key_size is None:
            alg_id = u32be_len_input(headers['enc'])
        else:
            alg_id = u32be_len_input(headers['alg'])

        # PartyUInfo
        apu_info = u32be_len_input(headers.get('apu'), True)

        # PartyVInfo
        apv_info = u32be_len_input(headers.get('apv'), True)

        # SuppPubInfo
        pub_info = struct.pack('>I', bit_size)

        return alg_id + apu_info + apv_info + pub_info

    def compute_derived_key(self, shared_key, fixed_info, bit_size):
        ckdf = ConcatKDFHash(
            algorithm=hashes.SHA256(),
            length=bit_size // 8,
            otherinfo=fixed_info,
            backend=default_backend()
        )
        return ckdf.derive(shared_key)

    def deliver(self, key, pubkey, headers, bit_size):
        shared_key = key.exchange_shared_key(pubkey)
        fixed_info = self.compute_fixed_info(headers, bit_size)
        return self.compute_derived_key(shared_key, fixed_info, bit_size)

    def _generate_ephemeral_key(self, key):
        return key.generate_key(key['crv'], is_private=True)

    def _prepare_headers(self, epk):
        # REQUIRED_JSON_FIELDS contains only public fields
        pub_epk = {k: epk[k] for k in epk.REQUIRED_JSON_FIELDS}
        pub_epk['kty'] = epk.kty
        return {'epk': pub_epk}

    def wrap(self, enc_alg, headers, key, preset=None):
        if self.key_size is None:
            bit_size = enc_alg.CEK_SIZE
        else:
            bit_size = self.key_size

        if preset and 'epk' in preset:
            epk = preset['epk']
            h = {}
        else:
            epk = self._generate_ephemeral_key(key)
            h = self._prepare_headers(epk)

        public_key = key.get_op_key('wrapKey')
        dk = self.deliver(epk, public_key, headers, bit_size)

        if self.key_size is None:
            return {'ek': b'', 'cek': dk, 'header': h}

        if preset and 'cek' in preset:
            preset_for_kw = {'cek': preset['cek']}
        else:
            preset_for_kw = None

        kek = self.aeskw.prepare_key(dk)
        rv = self.aeskw.wrap(enc_alg, headers, kek, preset_for_kw)
        rv['header'] = h
        return rv

    def unwrap(self, enc_alg, ek, headers, key):
        if 'epk' not in headers:
            raise ValueError('Missing "epk" in headers')

        if self.key_size is None:
            bit_size = enc_alg.CEK_SIZE
        else:
            bit_size = self.key_size

        epk = key.import_key(headers['epk'])
        public_key = epk.get_op_key('wrapKey')
        dk = self.deliver(key, public_key, headers, bit_size)

        if self.key_size is None:
            return dk

        kek = self.aeskw.prepare_key(dk)
        return self.aeskw.unwrap(enc_alg, ek, headers, kek)


def u32be_len_input(s, base64=False):
    if not s:
        return b'\x00\x00\x00\x00'
    if base64:
        s = urlsafe_b64decode(to_bytes(s))
    else:
        s = to_bytes(s)
    return struct.pack('>I', len(s)) + s


JWE_ALG_ALGORITHMS = [
    DirectAlgorithm(),  # dir
    RSAAlgorithm('RSA1_5', 'RSAES-PKCS1-v1_5', padding.PKCS1v15()),
    RSAAlgorithm(
        'RSA-OAEP', 'RSAES OAEP using default parameters',
        padding.OAEP(padding.MGF1(hashes.SHA1()), hashes.SHA1(), None)),
    RSAAlgorithm(
        'RSA-OAEP-256', 'RSAES OAEP using SHA-256 and MGF1 with SHA-256',
        padding.OAEP(padding.MGF1(hashes.SHA256()), hashes.SHA256(), None)),

    AESAlgorithm(128),  # A128KW
    AESAlgorithm(192),  # A192KW
    AESAlgorithm(256),  # A256KW
    AESGCMAlgorithm(128),  # A128GCMKW
    AESGCMAlgorithm(192),  # A192GCMKW
    AESGCMAlgorithm(256),  # A256GCMKW
    ECDHESAlgorithm(None),  # ECDH-ES
    ECDHESAlgorithm(128),  # ECDH-ES+A128KW
    ECDHESAlgorithm(192),  # ECDH-ES+A192KW
    ECDHESAlgorithm(256),  # ECDH-ES+A256KW
]

# 'PBES2-HS256+A128KW': '',
# 'PBES2-HS384+A192KW': '',
# 'PBES2-HS512+A256KW': '',
