use std::fmt;

use aes::Aes128;
use base64::{Engine, engine::general_purpose::STANDARD};
use cbc::{
    Decryptor, Encryptor,
    cipher::{BlockDecryptMut, BlockEncryptMut, KeyIvInit, block_padding::Pkcs7},
};

pub struct Cryptor {
    encryptor: Encryptor<Aes128>,
    decryptor: Decryptor<Aes128>,
}

impl Cryptor {
    pub fn new(key: &[u8], iv: &[u8]) -> Result<Self, cbc::cipher::InvalidLength> {
        Ok(Self {
            encryptor: Encryptor::new_from_slices(key, iv)?,
            decryptor: Decryptor::new_from_slices(key, iv)?,
        })
    }

    pub fn decrypt(&self, ciphertext: String) -> Result<String, DecryptError> {
        let mut ciphertext = STANDARD.decode(ciphertext).map_err(DecryptError::Base64)?;

        let plaintext = self
            .decryptor
            .clone()
            .decrypt_padded_mut::<Pkcs7>(ciphertext.as_mut_slice())
            .map_err(DecryptError::Unpad)?;

        Ok(String::from_utf8_lossy(plaintext).to_string())
    }

    pub fn encrypt(&self, plaintext: String) -> String {
        // Allocate buffer with extra space for padding (AES block size is 16 bytes)
        let plaintext_bytes = plaintext.as_bytes();
        let mut buffer = vec![0u8; 16 * (plaintext_bytes.len() / 16 + 1)];
        buffer[..plaintext_bytes.len()].copy_from_slice(plaintext_bytes);
        let ciphertext = self
            .encryptor
            .clone()
            .encrypt_padded_mut::<Pkcs7>(&mut buffer, plaintext_bytes.len())
            .expect("enough space for encrypting is allocated");

        STANDARD.encode(ciphertext)
    }
}

#[derive(Debug)]
pub enum DecryptError {
    Base64(base64::DecodeError),
    Unpad(cbc::cipher::block_padding::UnpadError),
}

impl fmt::Display for DecryptError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            DecryptError::Base64(e) => write!(f, "Base64 decode failed: {}", e),
            DecryptError::Unpad(e) => write!(f, "PKCS7 unpadding failed: {}", e),
        }
    }
}

impl std::error::Error for DecryptError {}