use anyhow::{Error, Result, anyhow};
use base64::prelude::*;
use crypto::{
    aes, blockmodes,
    buffer::{BufferResult, ReadBuffer, RefReadBuffer, RefWriteBuffer, WriteBuffer},
    symmetriccipher::SymmetricCipherError,
};
use rand::prelude::*;

pub const DEFAULT_KEY: &[u8] = "394e46c36e1844a88a4ea5804774d5ad".as_bytes();

fn gen_iv() -> [u8; 16] {
    let mut rng = rand::rng();
    let mut iv: [u8; 16] = [0; 16];
    rng.fill(&mut iv);
    iv
}

fn cipher_error(e: SymmetricCipherError) -> Error {
    match e {
        SymmetricCipherError::InvalidLength => anyhow!("InvalidLength"),
        SymmetricCipherError::InvalidPadding => anyhow!("InvalidPadding"),
    }
}

pub fn encrypt<T: AsRef<[u8]>>(key: &[u8], data: T) -> Result<String> {
    let iv = gen_iv();
    let mut encryptor =
        aes::cbc_encryptor(aes::KeySize::KeySize256, key, &iv, blockmodes::PkcsPadding);
    let mut final_result = Vec::from(iv);
    let mut read_buffer = RefReadBuffer::new(data.as_ref());
    let mut buffer = [0; 4096];
    let mut write_buffer = RefWriteBuffer::new(&mut buffer);

    loop {
        let result = encryptor
            .encrypt(&mut read_buffer, &mut write_buffer, true)
            .map_err(cipher_error)?;

        final_result.extend(
            write_buffer
                .take_read_buffer()
                .take_remaining()
                .iter()
                .copied(),
        );

        match result {
            BufferResult::BufferUnderflow => break,
            BufferResult::BufferOverflow => {}
        }
    }

    Ok(BASE64_URL_SAFE.encode(final_result))
}

pub fn decrypt<T: AsRef<[u8]>>(key: &[u8], data: T) -> Result<String> {
    let mut iv = BASE64_URL_SAFE.decode(data)?;
    let data = iv.split_off(16);

    let mut decryptor =
        aes::cbc_decryptor(aes::KeySize::KeySize256, key, &iv, blockmodes::PkcsPadding);
    let mut final_result = Vec::<u8>::new();
    let mut read_buffer = RefReadBuffer::new(data.as_ref());
    let mut buffer = [0; 4096];
    let mut write_buffer = RefWriteBuffer::new(&mut buffer);

    loop {
        let result = decryptor
            .decrypt(&mut read_buffer, &mut write_buffer, true)
            .map_err(cipher_error)?;

        final_result.extend(
            write_buffer
                .take_read_buffer()
                .take_remaining()
                .iter()
                .copied(),
        );

        match result {
            BufferResult::BufferUnderflow => break,
            BufferResult::BufferOverflow => {}
        }
    }

    Ok(String::from_utf8(final_result)?)
}

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

    #[test]
    fn test_encrypt() {
        let res = encrypt(DEFAULT_KEY, "123456");
        println!("result: {:?}", res);
    }

    #[test]
    fn test_decrypt() {
        let res = decrypt(DEFAULT_KEY, "zewf420I2cX6K3jtkUKB9n4W_3e2mi19uPYlNOxVT70=");
        println!("result: {:?}", res);
    }
}
