use ::kitty_pool::pool::{kitty_pool, KittyPoolOrganization};
use criterion::{criterion_group, criterion_main, Criterion};
use futures::executor::block_on;
use rand::{distributions::Alphanumeric, thread_rng, Rng};

fn test_con_fn(pool_size: usize, block_size: usize, request_size: usize) {
    let result = kitty_pool(KittyPoolOrganization::Contiguous, pool_size, block_size);
    assert!(result.is_ok());
    let mut kitty_pool = result.unwrap();
    let rand_string: String = thread_rng()
        .sample_iter(&Alphanumeric)
        .take(request_size)
        .collect();
    let rand_bytes: Vec<u8> = Vec::from(rand_string.as_bytes());
    let mut buffer = vec![0; request_size];
    let future = async {
        let token = kitty_pool.borrow(request_size).await;
        assert!(token.is_ok());
        let token = kitty_pool.write(&token.unwrap(), &rand_bytes).await;
        assert!(token.is_ok());
        let token = kitty_pool.read(&token.unwrap(), &mut buffer).await;
        assert!(token.is_ok());
        assert_eq!(*buffer, *rand_bytes);
        assert!(kitty_pool.release(&token.unwrap()).is_ok());
    };
    block_on(future);
}

fn small_con_benchmark(c: &mut Criterion) {
    const POOL_SIZE: usize = 1024 * 1024;
    const BLOCK_SIZE: usize = POOL_SIZE / 64;
    const REQUEST_SIZE: usize = BLOCK_SIZE * 4;
    c.bench_function("con small", |b| {
        b.iter(|| test_con_fn(POOL_SIZE, BLOCK_SIZE, REQUEST_SIZE))
    });
}

fn medium_con_benchmark(c: &mut Criterion) {
    const POOL_SIZE: usize = 1024 * 1024;
    const BLOCK_SIZE: usize = POOL_SIZE / 16;
    const REQUEST_SIZE: usize = BLOCK_SIZE * 2;
    c.bench_function("con medium", |b| {
        b.iter(|| test_con_fn(POOL_SIZE, BLOCK_SIZE, REQUEST_SIZE))
    });
}

fn large_con_benchmark(c: &mut Criterion) {
    const POOL_SIZE: usize = 1024 * 1024;
    const BLOCK_SIZE: usize = POOL_SIZE / 4;
    const REQUEST_SIZE: usize = BLOCK_SIZE * 1;
    c.bench_function("con large", |b| {
        b.iter(|| test_con_fn(POOL_SIZE, BLOCK_SIZE, REQUEST_SIZE))
    });
}

criterion_group!(
    con_benches,
    small_con_benchmark,
    medium_con_benchmark,
    large_con_benchmark
);

fn test_sgl_fn(pool_size: usize, block_size: usize, request_size: usize) {
    let result = kitty_pool(
        KittyPoolOrganization::ScatterGatherList,
        pool_size,
        block_size,
    );
    assert!(result.is_ok());
    let mut kitty_pool = result.unwrap();
    let rand_string: String = thread_rng()
        .sample_iter(&Alphanumeric)
        .take(request_size)
        .collect();
    let rand_bytes: Vec<u8> = Vec::from(rand_string.as_bytes());
    let mut buffer = vec![0; request_size];
    let future = async {
        let token = kitty_pool.borrow(request_size).await;
        assert!(token.is_ok());
        let token = kitty_pool.write(&token.unwrap(), &rand_bytes).await;
        assert!(token.is_ok());
        let token = kitty_pool.read(&token.unwrap(), &mut buffer).await;
        assert!(token.is_ok());
        assert_eq!(*buffer, *rand_bytes);
        assert!(kitty_pool.release(&token.unwrap()).is_ok());
    };
    block_on(future);
}

fn small_sgl_benchmark(c: &mut Criterion) {
    const POOL_SIZE: usize = 1024 * 1024;
    const BLOCK_SIZE: usize = POOL_SIZE / 64;
    const REQUEST_SIZE: usize = BLOCK_SIZE * 4;
    c.bench_function("sgl small", |b| {
        b.iter(|| test_sgl_fn(POOL_SIZE, BLOCK_SIZE, REQUEST_SIZE))
    });
}

fn medium_sgl_benchmark(c: &mut Criterion) {
    const POOL_SIZE: usize = 1024 * 1024;
    const BLOCK_SIZE: usize = POOL_SIZE / 16;
    const REQUEST_SIZE: usize = BLOCK_SIZE * 2;
    c.bench_function("sgl medium", |b| {
        b.iter(|| test_sgl_fn(POOL_SIZE, BLOCK_SIZE, REQUEST_SIZE))
    });
}

fn large_sgl_benchmark(c: &mut Criterion) {
    const POOL_SIZE: usize = 1024 * 1024;
    const BLOCK_SIZE: usize = POOL_SIZE / 4;
    const REQUEST_SIZE: usize = BLOCK_SIZE * 1;
    c.bench_function("sgl large", |b| {
        b.iter(|| test_sgl_fn(POOL_SIZE, BLOCK_SIZE, REQUEST_SIZE))
    });
}

criterion_group!(
    sgl_benches,
    small_sgl_benchmark,
    medium_sgl_benchmark,
    large_sgl_benchmark
);

criterion_main!(con_benches, sgl_benches);