use std::cmp::min;

use crate::common::{paddings, PaddingMethod};

use super::constants::{SM4_CK as CK, SM4_FK as FK, SM4_SBOX as SBOX, SM4_BLOCK_SIZE as BLOCK_SIZE};
use super::{operation_modes, Operation, OperationMode, BlockCipherImpl};


/// SM4分组密码算法流式计算器（SM4 Block Cipher Algorthm (Stream Mode) Calculator）
///
/// # 示例（Examples）
///
/// * 采用CBC工作模式，填充方法为第1套进行加密。
/// ```rust
/// # fn main() -> Result<(), ()> {
/// use neuedu_cryptos::common::PaddingMethod;
/// use neuedu_cryptos::block_ciphers::{BlockCipherImpl, Operation, OperationMode, SM4};
///
/// const KEY: [u8; 16] = [
///     0x01, 0x23, 0x45, 0x67, 0x89, 0xAB, 0xCD, 0xEF,
///     0xFE, 0xDC, 0xBA, 0x98, 0x76, 0x54, 0x32, 0x10,
/// ];
/// const IV: [u8; 16] = [
///     0x01, 0x23, 0x45, 0x67, 0x89, 0xAB, 0xCD, 0xEF,
///     0xFE, 0xDC, 0xBA, 0x98, 0x76, 0x54, 0x32, 0x10,
/// ];
/// const PLAIN: [u8; 16] = [
///     0x01, 0x23, 0x45, 0x67, 0x89, 0xAB, 0xCD, 0xEF,
///     0xFE, 0xDC, 0xBA, 0x98, 0x76, 0x54, 0x32, 0x10,
/// ];
///
/// let mut instance = SM4::new(
///     KEY,
///     Operation::Encrypt,
///     Some(OperationMode::CBC(IV.into())),
///     Some(PaddingMethod::Method1),
/// )?;
///
/// let mut cipher = Vec::<u8>::new();
/// cipher.append(&mut instance.update(&PLAIN)?);
/// cipher.append(&mut instance.r#final()?);
///
/// assert_eq!(cipher, [
///     0x26, 0x77, 0xF4, 0x6B, 0x09, 0xC1, 0x22, 0xCC,
///     0x97, 0x55, 0x33, 0x10, 0x5B, 0xD4, 0xA2, 0x2A,
///     0x3B, 0x88, 0x0E, 0x68, 0x67, 0x77, 0x25, 0x22,
///     0xAE, 0x55, 0xD2, 0xF0, 0xAE, 0x74, 0x78, 0xAE,
/// ]);
/// #
/// #     Ok(())
/// # }
/// ```
///
/// * 采用CBC工作模式，填充方法为第1套进行解密。
/// ```rust
/// # fn main() -> Result<(), ()> {
/// use neuedu_cryptos::common::PaddingMethod;
/// use neuedu_cryptos::block_ciphers::{BlockCipherImpl, Operation, OperationMode, SM4};
///
/// const KEY: [u8; 16] = [
///     0x01, 0x23, 0x45, 0x67, 0x89, 0xAB, 0xCD, 0xEF,
///     0xFE, 0xDC, 0xBA, 0x98, 0x76, 0x54, 0x32, 0x10,
/// ];
/// const IV: [u8; 16] = [
///     0x01, 0x23, 0x45, 0x67, 0x89, 0xAB, 0xCD, 0xEF,
///     0xFE, 0xDC, 0xBA, 0x98, 0x76, 0x54, 0x32, 0x10,
/// ];
/// const CIPHER: [u8; 32] = [
///     0x26, 0x77, 0xF4, 0x6B, 0x09, 0xC1, 0x22, 0xCC,
///     0x97, 0x55, 0x33, 0x10, 0x5B, 0xD4, 0xA2, 0x2A,
///     0x3B, 0x88, 0x0E, 0x68, 0x67, 0x77, 0x25, 0x22,
///     0xAE, 0x55, 0xD2, 0xF0, 0xAE, 0x74, 0x78, 0xAE,
/// ];
///
/// // 流式计算。
/// let mut instance = SM4::new(
///     KEY,
///     Operation::Decrypt,
///     Some(OperationMode::CBC(IV.into())),
///     Some(PaddingMethod::Method1),
/// )?;
///
/// let mut plain = Vec::<u8>::new();
/// plain.append(&mut instance.update(&CIPHER)?);
/// plain.append(&mut instance.r#final()?);
///
/// assert_eq!(plain, [
///     0x01, 0x23, 0x45, 0x67, 0x89, 0xAB, 0xCD, 0xEF,
///     0xFE, 0xDC, 0xBA, 0x98, 0x76, 0x54, 0x32, 0x10,
/// ]);
/// #
/// #     Ok(())
/// # }
/// ```
pub struct SM4 {
    mk: [u32; 4],
    rk: [u32; 32],

    operation: Operation,
    operation_mode: Option<Box<dyn operation_modes::OperationModeImpl>>,
    padding_method: Option<Box<dyn paddings::PaddingModeImpl>>,

    encrypt_buffer: [u8; BLOCK_SIZE as usize],
    encrypt_buffer_length: u8,

    decrypt_buffer: [u8; BLOCK_SIZE as usize],

    decrypt_expect_bytes: Option<u128>,
    received_bytes: u128,

    zero_padding_first_round: bool,
}

impl SM4 {

    /// 初始化一个SM4分组密码算法实例。\
    /// Initialize the SM4 block cipher algorithm instance.
    ///
    /// # 参数（Arguments）
    ///
    /// * key - 密钥（Key）
    ///
    ///   长度为16字节（128比特）的字节数组。\
    ///   A byte array of length 16 bytes (128 bits).
    ///
    /// * operation - 操作模式（Operating mode）
    ///
    ///   指示实例运行在加密模式或解密模式。\
    ///   Indicates whether the instance is running in encryption or decryption mode.
    ///
    /// * operation_mode - 分组工作模式（Mode of operation）
    ///
    ///   指示实例将使用的分组工作模式，如不进行分组，可以传入`None`。\
    ///   Indicates the mode of operation for the instance will use. If do not need, pass `None`.
    ///
    /// * padding_method - 填充方法（Method of padding）
    ///
    ///   指示实例将使用的填充方法，如不进行填充，可以传入`None`。\
    ///   Indicates the method of padding for the instance will use. If do not need, pass `None`.
    pub fn new(key: [u8; BLOCK_SIZE as usize], operation: Operation, operation_mode: Option<OperationMode>, padding_method: Option<PaddingMethod>) -> Result<Self, ()> {
        let operation_mode: Option<Box<dyn operation_modes::OperationModeImpl>> = match operation_mode {
            Some(operation_mode) => match operation_mode {
                OperationMode::ECB => Some(Box::<operation_modes::ECB>::default()),
                OperationMode::CBC(iv) => Some(Box::new(operation_modes::CBC::new(iv))),
                OperationMode::CFB(iv) => Some(Box::new(operation_modes::CFB::new(iv))),
                OperationMode::OFB(iv) => Some(Box::new(operation_modes::OFB::new(iv))),
            },
            None => None,
        };
        let padding_method: Option<Box<dyn paddings::PaddingModeImpl>> = match padding_method {
            Some(padding_mode) => {
                match padding_mode {
                    PaddingMethod::Method1 => Some(Box::new(paddings::Method1::new(BLOCK_SIZE as usize))),
                    PaddingMethod::Method2 => Some(Box::new(paddings::Method2::new(BLOCK_SIZE as usize))),
                    PaddingMethod::Method3(expect_sizes) => Some(Box::new(paddings::Method3::new(BLOCK_SIZE as usize, expect_sizes))),
                }
            },
            None => None,
        };
        if operation_mode.is_none() && padding_method.is_some() {
            return Err(());
        }

        let mk = [
            u32::from_be_bytes(key[0..4].try_into().unwrap()),
            u32::from_be_bytes(key[4..8].try_into().unwrap()),
            u32::from_be_bytes(key[8..12].try_into().unwrap()),
            u32::from_be_bytes(key[12..16].try_into().unwrap()),
        ];

        let mut instance = Self {
            mk,
            rk: [0u32; 32],

            operation,
            operation_mode,
            padding_method,

            encrypt_buffer: [0u8; BLOCK_SIZE as usize],
            encrypt_buffer_length: 0,

            decrypt_buffer: [0u8; BLOCK_SIZE as usize],
            decrypt_expect_bytes: None,

            received_bytes: 0u128,

            zero_padding_first_round: true,
        };

        let mut k = [0u32; 36];
        (k[0], k[1], k[2], k[3]) = instance.f6();
        for (i, item) in instance.rk.iter_mut().enumerate() {
            k[i + 4] = Self::f7(k[i], k[i + 1], k[i + 2], k[i + 3], CK[i]);
            *item = k[i + 4];
        }

        Ok(instance)
    }

    fn pre_update(&mut self) -> Box<[u8]> {
        if let Some(ref mut padding_mode) = self.padding_method {
            if let PaddingMethod::Method3(Some(expect_sizes)) = padding_mode.padding_mode() {
                if self.operation == Operation::Encrypt && self.zero_padding_first_round {
                    self.zero_padding_first_round = false;

                    let result = self.update(&(expect_sizes as u128 * 8).to_be_bytes()).unwrap();
                    self.received_bytes -= BLOCK_SIZE as u128;

                    return result.into();
                }
            }
        }

        Box::new([])
    }

    fn post_update(&mut self, result: &[u8]) -> Box<[u8]> {
        if let Some(ref mut padding_mode) = self.padding_method {
            if let PaddingMethod::Method3(_expect_sizes) = padding_mode.padding_mode() {
                if self.operation == Operation::Decrypt && self.zero_padding_first_round {
                    self.zero_padding_first_round = false;

                    self.decrypt_expect_bytes = Some(<u128>::from_be_bytes(result[..BLOCK_SIZE as usize].try_into().unwrap()) / 8);
                    self.received_bytes -= BLOCK_SIZE as u128;

                    return result[BLOCK_SIZE as usize..].into();
                }
            }
        }

        result.into()
    }

    fn calculate(&self, message: &[u8]) -> Box<[u8]> {
        let mut x = [0u32; 36];
        x[0..4].copy_from_slice(&[
            u32::from_be_bytes(message[0..4].try_into().unwrap()),
            u32::from_be_bytes(message[4..8].try_into().unwrap()),
            u32::from_be_bytes(message[8..12].try_into().unwrap()),
            u32::from_be_bytes(message[12..16].try_into().unwrap()),
        ]);

        for (i, _) in self.rk.iter().enumerate() {
            if self.operation == Operation::Encrypt {
                x[i + 4] = Self::f4(x[i], x[i + 1], x[i + 2], x[i + 3], self.rk[i]);

                continue;
            }

            if let Some(operation_mode) = &self.operation_mode {
                if operation_mode.decrypt_operation() == Operation::Encrypt {
                    x[i + 4] = Self::f4(x[i], x[i + 1], x[i + 2], x[i + 3], self.rk[i]);

                    continue;
                }
            }

            x[i + 4] = Self::f4(x[i], x[i + 1], x[i + 2], x[i + 3], self.rk[self.rk.len() - 1 - i]);
        }

        let y: [u32; 4] = Self::f5(x[32], x[33], x[34], x[35]).try_into().unwrap();

        [
            y[0].to_be_bytes().as_slice(),
            y[1].to_be_bytes().as_slice(),
            y[2].to_be_bytes().as_slice(),
            y[3].to_be_bytes().as_slice(),
        ].concat().into()
    }

    const fn f1(x0: u32, x1: u32, x2: u32, x3: u32, rk: u32) -> u32 {
        x0 ^ Self::f3(Self::f2(x1 ^ x2 ^ x3 ^ rk))
    }

    const fn f2(a: u32) -> u32 {
        let [a0, a1, a2, a3] = a.to_be_bytes();

        u32::from_be_bytes([SBOX[a0 as usize], SBOX[a1 as usize], SBOX[a2 as usize], SBOX[a3 as usize]])
    }

    const fn f3(b: u32) -> u32 {
        b ^ b.rotate_left(2) ^ b.rotate_left(10) ^ b.rotate_left(18) ^ b.rotate_left(24)
    }

    const fn f4(x0: u32, x1: u32, x2: u32, x3: u32, rk: u32) -> u32 {
        Self::f1(x0, x1, x2, x3, rk)
    }

    const fn f5(x32: u32, x33: u32, x34: u32, x35: u32) -> (u32, u32, u32, u32) {
        (x35, x34, x33, x32)
    }

    const fn f6(&self) -> (u32, u32, u32, u32) {
        (self.mk[0] ^ FK[0], self.mk[1] ^ FK[1], self.mk[2] ^ FK[2], self.mk[3] ^ FK[3])
    }

    const fn f7(k0: u32, k1: u32, k2: u32, k3: u32, ck: u32) -> u32 {
        k0 ^ Self::f8(Self::f2(k1 ^ k2 ^ k3 ^ ck))
    }

    const fn f8(b: u32) -> u32 {
        b ^ b.rotate_left(13) ^ b.rotate_left(23)
    }
}

impl BlockCipherImpl for SM4 {

    fn update(&mut self, message: &[u8]) -> Result<Vec<u8>, ()> {
        self.received_bytes += message.len() as u128;

        if self.operation_mode.is_none() && self.received_bytes > BLOCK_SIZE as u128 {
            return Err(())
        }

        let mut result = Vec::<u8>::with_capacity(message.len() + BLOCK_SIZE as usize * 2);

        result.append(&mut self.pre_update().into());

        let mut offset = 0usize;
        while offset < message.len() {
            let current_length = min(
                message.len() - offset,                             // 剩余消息长度
                (BLOCK_SIZE - self.encrypt_buffer_length) as usize  // 剩余缓冲区长度
            ) as u8;

            self.encrypt_buffer[self.encrypt_buffer_length as usize..self.encrypt_buffer_length as usize + current_length as usize].copy_from_slice(&message[offset..offset + current_length as usize]);
            offset += current_length as usize;
            self.encrypt_buffer_length += current_length;

            if self.encrypt_buffer_length == BLOCK_SIZE {
                let mut message = message.to_vec();
                if let Some(ref mut operation_mode) = &mut self.operation_mode {
                    message = operation_mode.pre_calculate(&self.encrypt_buffer, &self.operation).into();
                }
                message = self.calculate(&message).into();
                if let Some(ref mut operation_mode) = &mut self.operation_mode {
                    message = operation_mode.post_calculate(&message, &self.operation).into();
                }
                result.append(&mut message);

                self.encrypt_buffer_length = 0;
            }
        }

        result = self.post_update(&result).into();

        Ok(if result.is_empty() {
            vec![]
        } else {
            match self.operation {
                Operation::Encrypt => result,
                Operation::Decrypt => {
                    match &self.padding_method {
                        Some(_) => {
                            self.decrypt_buffer[..16].copy_from_slice(&result[result.len() - BLOCK_SIZE as usize..]);

                            result[0..result.len() - BLOCK_SIZE as usize].into()
                        },
                        None => result,
                    }
                },
            }
        })
    }

    fn r#final(&mut self) -> Result<Vec<u8>, ()> {
        // 如果不进行分组，并且接收到的数据长度不等于块长度，进行错误响应。
        if self.operation_mode.is_none() && self.received_bytes != BLOCK_SIZE as u128 {
            return Err(());
        }

        match &self.padding_method {
            Some(padding_mode) => match self.operation {
                Operation::Encrypt => {
                    // 如果选择填充方法3，将验证接收到的数据长度与预期长度是否一致。
                    if let PaddingMethod::Method3(Some(expect_sizes)) = padding_mode.padding_mode() {
                        if self.received_bytes != expect_sizes as u128 {
                            return Err(());
                        }
                    }

                    Ok(self.update(&padding_mode.pad(&self.encrypt_buffer[..self.encrypt_buffer_length as usize])).unwrap())
                },
                Operation::Decrypt => {
                    // 如果选择填充方法3，将响应结果调整至预期长度。
                    if let PaddingMethod::Method3(None) = padding_mode.padding_mode() {
                        let result = padding_mode.unpad(&self.decrypt_buffer);

                        return Ok(result[..result.len() - (self.received_bytes - self.decrypt_expect_bytes.expect("")) as usize].into())
                    }

                    Ok(padding_mode.unpad(&self.decrypt_buffer).into())
                },
            },
            None => Ok(vec![]),
        }
    }
}
