use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};

use leona::{KeyedHash, Lioness, ZeroIv};

use aes::Aes256;
use blake2::Blake2sMac256;
use chacha20::ChaCha20;
use ctr::Ctr32BE;
use hmac::Hmac;
use sha2::Sha256;
use sha3::Sha3_256;
use typenum::U32;

use cipher::stream::StreamCipher;
use crypto_common::{KeyInit, KeySizeUser};
use digest::{FixedOutput, Update};
use hybrid_array::Array;
use typenum::{IsLessOrEqual, True, U64};
use zerocopy::FromZeros;

fn inner<S, H>(c: &mut Criterion, name_cipher: &str, name_hash: &str)
where
    S: KeySizeUser<KeySize = H::OutputSize> + KeyInit + StreamCipher,
    H: KeySizeUser + KeyInit + FixedOutput + Update,
    S::KeySize: IsLessOrEqual<U64, Output = True>,
    H::KeySize: IsLessOrEqual<U64, Output = True>,
    Array<u8, S::KeySize>: FromZeros,
{
    const KB: usize = 1024;

    let mut group = c.benchmark_group(format!("{name_cipher}/{name_hash}"));
    let cipher = Lioness::<S, H>::new_dynamic(b"");

    for size in [KB, 2 * KB, 4 * KB, 8 * KB, 16 * KB] {
        let mut buf = vec![0u8; size];

        group.throughput(Throughput::Bytes(size as u64));

        group.bench_function(BenchmarkId::new("encrypt", size), |b| {
            b.iter(|| cipher.encrypt(&mut buf));
        });
    }
}

fn bench(c: &mut Criterion) {
    inner::<ZeroIv<ChaCha20>, Blake2sMac256>(c, "chacha20", "blake2s256");
    inner::<ZeroIv<ChaCha20>, KeyedHash<U32, Sha256>>(c, "chacha20", "sha256");
    inner::<ZeroIv<ChaCha20>, KeyedHash<U32, Sha3_256>>(c, "chacha20", "sha3_256");
    inner::<ZeroIv<ChaCha20>, Hmac<Sha256>>(c, "chacha20", "sha256-hmac");
    inner::<ZeroIv<Ctr32BE<Aes256>>, Blake2sMac256>(c, "aes256", "blake2s256");
    inner::<ZeroIv<Ctr32BE<Aes256>>, KeyedHash<U32, Sha256>>(c, "aes256", "sha256");
    inner::<ZeroIv<Ctr32BE<Aes256>>, KeyedHash<U32, Sha3_256>>(c, "aes256", "sha3_256");
    inner::<ZeroIv<Ctr32BE<Aes256>>, Hmac<Sha256>>(c, "aes256", "sha256-hmac");
}

criterion_group!(benches, bench);
criterion_main!(benches);