use crate::error::DatError;
use aes_gcm::aead::array::Array;
use aes_gcm::aead::common::Generate;
use aes_gcm::aead::consts::{U12, U16, U32};
use aes_gcm::aead::Aead;
use aes_gcm::{Aes128Gcm, AesGcm};
use aes_gcm::{Aes256Gcm, Key, KeyInit, Nonce};
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum CryptAlgorithm {
    AES128GCMN,
    AES256GCMN,
}

impl CryptAlgorithm {
    pub fn to_str(&self) -> &'static str {
        match self {
            CryptAlgorithm::AES128GCMN => "AES128GCMN",
            CryptAlgorithm::AES256GCMN => "AES256GCMN",
        }
    }
    pub fn from_str(s: &str) -> Result<CryptAlgorithm, DatError> {
        match s {
            "AES128GCMN" => Ok(CryptAlgorithm::AES128GCMN),
            "AES256GCMN" => Ok(CryptAlgorithm::AES256GCMN),
            _ => Err(DatError::CryptError(format!("UnknownCryptAlgorithm: {}", s))),
        }
    }
}

pub enum CryptKey {
    AES128GCMN(Array<u8, U16>),
    AES256GCMN(Array<u8, U32>),
}

pub enum Cipher {
    AES128GCMN(AesGcm<aes_gcm::aes::Aes128, U12>),
    AES256GCMN(AesGcm<aes_gcm::aes::Aes256, U12>),
}

impl CryptKey {
    pub fn to_alg(&self) -> CryptAlgorithm {
        match self {
            CryptKey::AES128GCMN(_) => CryptAlgorithm::AES128GCMN,
            CryptKey::AES256GCMN(_) => CryptAlgorithm::AES256GCMN,
        }
    }
    pub fn from_bytes(alg: CryptAlgorithm, data: &[u8]) -> Result<CryptKey, DatError> {
        Ok(match alg {
            CryptAlgorithm::AES128GCMN => CryptKey::AES128GCMN(
                Array::<u8, U16>::try_from(data)
                    .map_err(|e| DatError::CryptError(format!("invalid crypt key: {}: {}", alg.to_str(), e)))?
            ),
            CryptAlgorithm::AES256GCMN => CryptKey::AES256GCMN(
                Array::<u8, U32>::try_from(data)
                    .map_err(|e| DatError::CryptError(format!("invalid crypt key: {}: {}", alg.to_str(), e)))?
            ),
        })
    }
    pub fn to_bytes(&self) -> Box<[u8]> {
        match self {
            CryptKey::AES128GCMN(c) => Box::from(c.as_slice()),
            CryptKey::AES256GCMN(c) => Box::from(c.as_slice()),
        }
    }
    pub fn to_cipher(&self) -> Cipher {
        match self {
            CryptKey::AES128GCMN(c) => Cipher::AES128GCMN(Aes128Gcm::new(c)),
            CryptKey::AES256GCMN(c) => Cipher::AES256GCMN(Aes256Gcm::new(c)),
        }
    }
    pub fn generate(alg: CryptAlgorithm) -> Self {
        match alg {
            CryptAlgorithm::AES128GCMN => CryptKey::AES128GCMN(Key::<Aes128Gcm>::generate()),
            CryptAlgorithm::AES256GCMN => CryptKey::AES256GCMN(Key::<Aes256Gcm>::generate()),
        }
    }
}

impl Cipher {
    pub fn encrypt(&self, body: &[u8]) -> Result<Vec<u8>, DatError> {
        if body.is_empty() {
            return Ok(vec![]);
        }

        macro_rules! aes_gcmn_encrypt {
            ($cipher:expr) => {{
                let nonce = Nonce::generate();
                let mut enc_data = nonce.to_vec();
                enc_data.extend(
                    $cipher.encrypt(&nonce, body)
                        .map_err(|e| DatError::CryptError(format!("encrypt error {}", e)))?
                );
                Ok(enc_data)
            }};
        }
        match self {
            Cipher::AES128GCMN(c) => aes_gcmn_encrypt!(c),
            Cipher::AES256GCMN(c) => aes_gcmn_encrypt!(c),
        }
    }
    pub fn decrypt(&self, data: &[u8]) -> Result<Vec<u8>, DatError> {
        if data.is_empty() {
            return Ok(vec![]);
        }

        macro_rules! aes_gcmn_decrypt {
            ($cipher:expr) => {{
                let nonce = data[..12].try_into()
                    .map_err(|_| DatError::CryptError(format!("decrypt nonce error")))?;

                $cipher.decrypt(&nonce, &data[12..])
                    .map_err(|e| DatError::CryptError(format!("decrypt error {}", e)))
            }};
        }
        match self {
            Cipher::AES128GCMN(c) => aes_gcmn_decrypt!(c),
            Cipher::AES256GCMN(c) => aes_gcmn_decrypt!(c),
        }
    }
}
