use std::sync::atomic::{AtomicU16, Ordering::SeqCst};

use aes::{
    cipher::{
        block_padding::NoPadding, generic_array::GenericArray, BlockDecrypt, BlockDecryptMut,
        BlockEncrypt, BlockEncryptMut, KeyInit, KeyIvInit,
    },
    Aes128,
};
use anyhow::{bail, Result};
use rand::{self, RngCore};
use sha2::{Digest, Sha256};

use crate::hex;

#[allow(non_camel_case_types)]
#[repr(u8)]
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)]
pub enum MsgType {
    HANDSHAKE_REQUEST = 0x0,
    ENCRYPTED_RESPONSE = 0x3,
    ENCRYPTED_REQUEST = 0x6,
}

#[derive(Debug, Default)]
pub struct Security {
    request_count: AtomicU16,
    response_count: AtomicU16,

    tcp_key: Option<[u8; 32]>,
}

impl Security {
    const N: u128 = 141661095494369103254425781617665632877;
    const KEY: [u8; 16] = Self::N.to_be_bytes();
    const IV: [u8; 16] = [b'\0'; 16];
    const SALT: &[u8] = b"xhdiwjnchekd4d512chdjx5d8e4c394D2D7S";

    pub fn encode32_data(data: &[u8]) -> Vec<u8> {
        md5::compute([data, Self::SALT].concat()).to_vec()
    }

    pub fn aes_encrypt(data: &mut [u8]) -> Vec<u8> {
        let array = GenericArray::from(Self::KEY);
        let cipher = Aes128::new(&array);

        let mut result = vec![0; data.len()];
        let padding = 16 - (result.len() % 16);
        result.extend(&vec![0; padding]);

        for (inchunk, outchunk) in data.chunks(16).zip(result.chunks_mut(16)) {
            let in_data = if inchunk.len() != 16 {
                [inchunk, &vec![0; padding]].concat()
            } else {
                inchunk.to_vec()
            };

            let inblock = GenericArray::from_slice(&in_data);
            let mut out_block = GenericArray::from_mut_slice(outchunk);

            cipher.encrypt_block_b2b(&inblock, &mut out_block);
        }

        result
    }

    pub fn aes_decrypt(data: &mut [u8]) {
        let array = GenericArray::from(Self::KEY);
        let cipher = Aes128::new(&array);

        for chunk in data.chunks_mut(16) {
            let mut block = GenericArray::from_mut_slice(chunk);
            cipher.decrypt_block(&mut block);
        }
    }

    pub fn aes_cbc_encrypt(&self, raw: &[u8], key: &[u8; 32]) -> Vec<u8> {
        type Aes256CbcEnc = cbc::Encryptor<aes::Aes256>;

        debug_assert_eq!(raw.len() % key.len(), 0);

        raw.chunks(key.len())
            .map(|r| {
                Aes256CbcEnc::new(key.into(), &Self::IV.into())
                    .encrypt_padded_vec_mut::<NoPadding>(r)
            })
            .flatten()
            .collect()
    }

    pub fn aes_cbc_decrypt(&self, raw: [u8; 32], key: &[u8; 32]) -> [u8; 32] {
        type Aes256CbcDec = cbc::Decryptor<aes::Aes256>;

        Aes256CbcDec::new(key.into(), &Self::IV.into())
            .decrypt_padded_vec_mut::<NoPadding>(&raw)
            .unwrap()
            .try_into()
            .unwrap()
    }

    pub fn tcp_key(&mut self, response: &[u8], key: &[u8; 32]) -> Result<()> {
        if response == b"ERROR" {
            bail!("authentication failed! (code ERROR)");
        }

        let payload: [u8; 32] = response[0..32]
            .iter()
            .map(|&b| b)
            .collect::<Vec<u8>>()
            .try_into()
            .unwrap();

        let sign = &response[32..];
        let result = self.aes_cbc_decrypt(payload, &key);

        if Sha256::digest(&result).into_iter().collect::<Vec<u8>>() != sign {
            bail!("sign does not match");
        }

        self.tcp_key = Some(Self::xorstr(&result, key).try_into().unwrap());
        self.request_count.store(0, SeqCst);
        self.response_count.store(0, SeqCst);

        Ok(())
    }

    fn xorstr(lhs: &[u8], rhs: &[u8]) -> Vec<u8> {
        debug_assert_eq!(lhs.len(), rhs.len());

        lhs.iter().zip(rhs.iter()).map(|(&l, &r)| l ^ r).collect()
    }

    pub fn encode_8370(&self, msg: &[u8], msg_type: MsgType) -> Result<Vec<u8>> {
        let mut header = hex("8370")?;
        let mut data: Vec<u8> = msg.to_vec();

        let mut size = data.len() as u16;
        let mut padding = 0;

        if msg_type == MsgType::ENCRYPTED_RESPONSE || msg_type == MsgType::ENCRYPTED_REQUEST {
            if (size + 2) % 16 != 0 {
                padding = 16 - ((size + 2) & 0xf);
                size += padding + 32;
                data.extend({
                    let mut d = vec![0; padding as usize];
                    rand::thread_rng().fill_bytes(&mut d);

                    d
                });
            }
        }

        header.extend(size.to_be_bytes());
        header.extend([0x20, (padding << 4) as u8 | msg_type as u8]);

        data = {
            let mut b = self
                .request_count
                .fetch_add(1, SeqCst)
                .to_be_bytes()
                .to_vec();
            b.extend(data);
            b
        };

        if msg_type == MsgType::ENCRYPTED_RESPONSE || msg_type == MsgType::ENCRYPTED_REQUEST {
            let sign: Vec<u8> = Sha256::digest(Self::add_bytes(header.clone(), data.clone()))
                .into_iter()
                .collect();

            data = self.aes_cbc_encrypt(&data, self.tcp_key.as_ref().unwrap());
            data.extend(sign);
        }

        header.extend(data);
        Ok(header)
    }

    fn add_bytes(mut lhs: Vec<u8>, rhs: Vec<u8>) -> Vec<u8> {
        lhs.extend(rhs);
        lhs
    }

    pub fn decode_8370(&self, mut data: Vec<u8>) -> Result<(Vec<Vec<u8>>, Vec<u8>)> {
        if data.len() < 6 {
            return Ok((Vec::new(), data));
        }

        let header = data[..6].to_vec();

        if header[0] != 0x83 || header[1] != 0x70 {
            bail!("no an 8370 message");
        } else if header[4] != 0x20 {
            bail!("missing byte 4");
        }

        let size = u16::from_be_bytes(header[2..4].try_into().unwrap()) as usize + 8;

        let mut leftover = Vec::new();

        if data.len() < size {
            return Ok((Vec::new(), data.to_vec()));
        } else if data.len() > size {
            leftover = data[size..].to_vec();
            data = data[..size].to_vec();
        }

        let padding = header[5] >> 4;
        let msgtype = header[5] & 0xf;
        data = data[6..].to_vec();

        if msgtype == MsgType::ENCRYPTED_RESPONSE as u8
            || msgtype == MsgType::ENCRYPTED_REQUEST as u8
        {
            let sign = data[(data.len() - 32)..].to_vec();
            data = data[..(data.len() - 32)].to_vec();
            data = self
                .aes_cbc_decrypt(data.try_into().unwrap(), &self.tcp_key.unwrap())
                .to_vec();

            let compare_sign: Vec<u8> = Sha256::digest(&[header.to_vec(), data.clone()].concat())
                .into_iter()
                .collect();

            if sign != compare_sign {
                bail!("sign does not match");
            }

            if padding != 0 {
                data = data[..(data.len() - padding as usize)].to_vec();
            }
        }

        self.response_count
            .store(u16::from_be_bytes(data[..2].try_into().unwrap()), SeqCst);

        data = data[2..].to_vec();

        if leftover.len() > 0 {
            let (mut packets, incomplete) = self.decode_8370(leftover)?;

            packets.insert(0, data);

            Ok((packets, incomplete))
        } else {
            Ok((vec![data], Vec::new()))
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn aes_cbc_decrypt() {
        let payload: [u8; 32] =
            b",\xcbq_T\x81L\x96\xfa\xe7\xe4\xa7\xc5\xabU \r\xf5x\xd6\x08\x94_\\\xce\x8br\x1b\xa5\xbe\xc6\x1a"
            .iter()
            .map(|&b| b)
            .collect::<Vec<u8>>()
            .try_into()
            .unwrap();

        let key = b"*[R\x00\xc2\xc0ML\x81\x1d\x05P\xe1\xdc[1CT6\xb9[wM*\x88\xd7\xe4ma\xfd\x96i";
        let plain = b"\x9b\xaa\xdf\xff\x07\x1a\xd2\xe4\xb7TY\xe2\xf9\x8c\xdf\xe7!+\xda\xe4\x86GY\xe6j\x94\xdb\xe7\xb9b\xda\xe6";

        let security = Security::default();

        let result = security.aes_cbc_decrypt(payload, key);

        assert_eq!(&result, plain);
    }
}