use super::{

    arch::*,

    utils::{aesdec8, aesdeclast8, aesenc8, aesenclast8, load8, store8, xor8, U128x8},

};

use crate::{Block, ParBlocks};

use cipher::{

    consts::{U16, U32, U8},

    generic_array::GenericArray,

    BlockCipher, BlockDecrypt, BlockEncrypt, NewBlockCipher,

};

mod expand;

#[cfg(test)]

mod test_expand;

/// AES-256 round keys

type RoundKeys = [__m128i; 15];

/// AES-256 block cipher

#[derive(Clone)]

pub struct Aes256 {

    encrypt_keys: RoundKeys,

    decrypt_keys: RoundKeys,

}

impl Aes256 {

    #[inline(always)]

    pub(crate) fn encrypt8(&self, mut blocks: U128x8) -> U128x8 {

        #[inline]

        #[target_feature(enable = "aes")]

        unsafe fn aesni256_encrypt8(keys: &RoundKeys, blocks: &mut U128x8) {

            xor8(blocks, keys[0]);

            aesenc8(blocks, keys[1]);

            aesenc8(blocks, keys[2]);

            aesenc8(blocks, keys[3]);

            aesenc8(blocks, keys[4]);

            aesenc8(blocks, keys[5]);

            aesenc8(blocks, keys[6]);

            aesenc8(blocks, keys[7]);

            aesenc8(blocks, keys[8]);

            aesenc8(blocks, keys[9]);

            aesenc8(blocks, keys[10]);

            aesenc8(blocks, keys[11]);

            aesenc8(blocks, keys[12]);

            aesenc8(blocks, keys[13]);

            aesenclast8(blocks, keys[14]);

        }

        unsafe { aesni256_encrypt8(&self.encrypt_keys, &mut blocks) };

        blocks

    }

    #[inline(always)]

    pub(crate) fn encrypt(&self, block: __m128i) -> __m128i {

        #[inline]

        #[target_feature(enable = "aes")]

        unsafe fn aesni256_encrypt1(keys: &RoundKeys, mut block: __m128i) -> __m128i {

            block = _mm_xor_si128(block, keys[0]);

            block = _mm_aesenc_si128(block, keys[1]);

            block = _mm_aesenc_si128(block, keys[2]);

            block = _mm_aesenc_si128(block, keys[3]);

            block = _mm_aesenc_si128(block, keys[4]);

            block = _mm_aesenc_si128(block, keys[5]);

            block = _mm_aesenc_si128(block, keys[6]);

            block = _mm_aesenc_si128(block, keys[7]);

            block = _mm_aesenc_si128(block, keys[8]);

            block = _mm_aesenc_si128(block, keys[9]);

            block = _mm_aesenc_si128(block, keys[10]);

            block = _mm_aesenc_si128(block, keys[11]);

            block = _mm_aesenc_si128(block, keys[12]);

            block = _mm_aesenc_si128(block, keys[13]);

            _mm_aesenclast_si128(block, keys[14])

        }

        unsafe { aesni256_encrypt1(&self.encrypt_keys, block) }

    }

}

impl NewBlockCipher for Aes256 {

    type KeySize = U32;

    #[inline]

    fn new(key: &GenericArray<u8, U32>) -> Self {

        let key = unsafe { &*(key as *const _ as *const [u8; 32]) };

        let (encrypt_keys, decrypt_keys) = expand::expand(key);

        Self {

            encrypt_keys,

            decrypt_keys,

        }

    }

}

impl BlockCipher for Aes256 {

    type BlockSize = U16;

    type ParBlocks = U8;

}

impl BlockEncrypt for Aes256 {

    #[inline]

    fn encrypt_block(&self, block: &mut Block) {

        // Safety: `loadu` and `storeu` support unaligned access

        #[allow(clippy::cast_ptr_alignment)]

        unsafe {

            let b = _mm_loadu_si128(block.as_ptr() as *const __m128i);

            let b = self.encrypt(b);

            _mm_storeu_si128(block.as_mut_ptr() as *mut __m128i, b);

        }

    }

    #[inline]

    fn encrypt_par_blocks(&self, blocks: &mut ParBlocks) {

        let b = self.encrypt8(load8(blocks));

        store8(blocks, b);

    }

}

impl BlockDecrypt for Aes256 {

    #[inline]

    fn decrypt_block(&self, block: &mut Block) {

        #[inline]

        #[target_feature(enable = "aes")]

        unsafe fn aes256_decrypt1(block: &mut Block, keys: &RoundKeys) {

            // Safety: `loadu` and `storeu` support unaligned access

            #[allow(clippy::cast_ptr_alignment)]

            let mut b = _mm_loadu_si128(block.as_ptr() as *const __m128i);

            b = _mm_xor_si128(b, keys[14]);

            b = _mm_aesdec_si128(b, keys[13]);

            b = _mm_aesdec_si128(b, keys[12]);

            b = _mm_aesdec_si128(b, keys[11]);

            b = _mm_aesdec_si128(b, keys[10]);

            b = _mm_aesdec_si128(b, keys[9]);

            b = _mm_aesdec_si128(b, keys[8]);

            b = _mm_aesdec_si128(b, keys[7]);

            b = _mm_aesdec_si128(b, keys[6]);

            b = _mm_aesdec_si128(b, keys[5]);

            b = _mm_aesdec_si128(b, keys[4]);

            b = _mm_aesdec_si128(b, keys[3]);

            b = _mm_aesdec_si128(b, keys[2]);

            b = _mm_aesdec_si128(b, keys[1]);

            b = _mm_aesdeclast_si128(b, keys[0]);

            // Safety: `loadu` and `storeu` support unaligned access

            #[allow(clippy::cast_ptr_alignment)]

            _mm_storeu_si128(block.as_mut_ptr() as *mut __m128i, b);

        }

        unsafe { aes256_decrypt1(block, &self.decrypt_keys) }

    }

    #[inline]

    fn decrypt_par_blocks(&self, blocks: &mut ParBlocks) {

        #[inline]

        #[target_feature(enable = "aes")]

        unsafe fn aes256_decrypt8(blocks: &mut ParBlocks, keys: &RoundKeys) {

            let mut b = load8(blocks);

            xor8(&mut b, keys[14]);

            aesdec8(&mut b, keys[13]);

            aesdec8(&mut b, keys[12]);

            aesdec8(&mut b, keys[11]);

            aesdec8(&mut b, keys[10]);

            aesdec8(&mut b, keys[9]);

            aesdec8(&mut b, keys[8]);

            aesdec8(&mut b, keys[7]);

            aesdec8(&mut b, keys[6]);

            aesdec8(&mut b, keys[5]);

            aesdec8(&mut b, keys[4]);

            aesdec8(&mut b, keys[3]);

            aesdec8(&mut b, keys[2]);

            aesdec8(&mut b, keys[1]);

            aesdeclast8(&mut b, keys[0]);

            store8(blocks, b);

        }

        unsafe { aes256_decrypt8(blocks, &self.decrypt_keys) }

    }

}

opaque_debug::implement!(Aes256);