use std::cmp::min;

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


pub struct SM4 {
    mk: [u32; 4],
    rk: [u32; 32],

    operation: Operation,
    operation_mode: Box<dyn operation_modes::OperationModeImpl>,
    padding_mode: 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 {

    pub fn new(key: [u8; BLOCK_SIZE as usize], operation: Operation, operation_mode: Option<OperationMode>, padding_mode: Option<PaddingMode>) -> Self {
        let mut instance = Self {
            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()),
            ],
            rk: [0u32; 32],

            operation,
            operation_mode: match operation_mode {
                Some(operation_mode) => match operation_mode {
                    OperationMode::ECB => Box::<operation_modes::ECB>::default(),
                    OperationMode::CBC(iv) => Box::new(operation_modes::CBC::new(iv)),
                    OperationMode::CFB(iv) => Box::new(operation_modes::CFB::new(iv)),
                    OperationMode::OFB(iv) => Box::new(operation_modes::OFB::new(iv)),
                },
                None => Box::<operation_modes::ECB>::default(),
            },
            padding_mode: match padding_mode {
                Some(padding_mode) => {
                    match padding_mode {
                        PaddingMode::Mode1 => Some(Box::<paddings::Mode1>::default()),
                        PaddingMode::Mode2 => Some(Box::<paddings::Mode2>::default()),
                        PaddingMode::Mode3(expect_sizes) => Some(Box::new(paddings::Mode3::new(expect_sizes))),
                    }
                },
                None => None,
            },

            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,
        };

        if let Some(ref mut padding_mode) = instance.padding_mode {
             (*padding_mode).set_block_size(BLOCK_SIZE as usize);
        }

        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];
        }

        instance
    }

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

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

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

        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 message = self.operation_mode.pre_calculate(&self.encrypt_buffer, &self.operation);
                let message = self.calculate(message);
                let mut message = self.operation_mode.post_calculate(&message, &self.operation);
                result.append(&mut message);

                self.encrypt_buffer_length = 0;
            }
        }

        result = self.post_update(result);

        match self.operation {
            Operation::Encrypt => result,
            Operation::Decrypt => {
                if result.is_empty() {
                    return vec![];
                }

                match &self.padding_mode {
                    Some(_) => {
                        self.decrypt_buffer[..16].copy_from_slice(&result[result.len() - BLOCK_SIZE as usize..]);

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

    pub fn r#final(&mut self) -> Vec<u8> {
        if let Some(ref padding_mode) = self.padding_mode {
            if let PaddingMode::Mode3(Some(expect_sizes)) = padding_mode.padding_mode() {
                if self.received_bytes > expect_sizes as u128 {
                    panic!("输入的长度大于预期！ The length of the input is greater than expected!");
                }

                return self.update(&padding_mode.pad(&self.encrypt_buffer[..self.encrypt_buffer_length as usize]));
            }

            if let PaddingMode::Mode3(None) = padding_mode.padding_mode() {
                let result = padding_mode.unpad(&self.decrypt_buffer);

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

            return match self.operation {
                Operation::Encrypt => self.update(&padding_mode.pad(&self.encrypt_buffer[..self.encrypt_buffer_length as usize])),
                Operation::Decrypt => padding_mode.unpad(&self.decrypt_buffer),
            };
        }

        vec![]
    }

    fn pre_update(&mut self) -> Vec<u8> {
        if let Some(ref mut padding_mode) = self.padding_mode {
            if let PaddingMode::Mode3(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).to_be_bytes());
                    self.received_bytes -= BLOCK_SIZE as u128;

                    return result;
                }
            }
        }

        vec![]
    }

    fn post_update(&mut self, result: Vec<u8>) -> Vec<u8> {
        if let Some(ref mut padding_mode) = self.padding_mode {
            if let PaddingMode::Mode3(_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()));
                    self.received_bytes -= BLOCK_SIZE as u128;

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

        result
    }

    fn calculate(&self, message: Vec<u8>) -> Vec<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() {
            x[i + 4] = if self.operation == Operation::Encrypt || self.operation_mode.decrypt_operation() == Operation::Encrypt {
                Self::f4(x[i], x[i + 1], x[i + 2], x[i + 3], self.rk[i])
            } else {
                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()
    }

    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)
    }
}
