use std::hint::black_box;
use std::thread;
use std::time::{Duration, Instant};

use criterion::async_executor::FuturesExecutor;
use criterion::{Criterion, Throughput, criterion_group, criterion_main};
use scc::HashMap;

fn insert_remove_single_async(c: &mut Criterion) {
    c.bench_function("HashMap: insert_remove_single_async", |b| {
        let hashmap: HashMap<u64, u64> = HashMap::default();
        assert!(hashmap.insert_sync(0, 0).is_ok());
        async fn test(hashmap: &HashMap<u64, u64>) {
            assert!(hashmap.insert_async(1, 1).await.is_ok());
            assert!(hashmap.remove_async(&1).await.is_some());
        }
        b.to_async(FuturesExecutor).iter(|| test(&hashmap));
    });
}

fn insert_remove_single_sync(c: &mut Criterion) {
    c.bench_function("HashMap: insert_remove_single_sync", |b| {
        let hashmap: HashMap<u64, u64> = HashMap::default();
        assert!(hashmap.insert_sync(0, 0).is_ok());
        fn test(hashmap: &HashMap<u64, u64>) {
            assert!(hashmap.insert_sync(1, 1).is_ok());
            assert!(hashmap.remove_sync(&1).is_some());
        }
        b.iter(|| test(&hashmap));
    });
}

fn insert_cold_async(c: &mut Criterion) {
    c.bench_function("HashMap: insert_cold_async", |b| {
        b.to_async(FuturesExecutor).iter_custom(async |iters| {
            let hashmap: HashMap<u64, u64> = HashMap::default();
            let start = Instant::now();
            for i in 0..iters {
                assert!(hashmap.insert_async(i, i).await.is_ok());
            }
            start.elapsed()
        })
    });
}

fn insert_cold_sync(c: &mut Criterion) {
    c.bench_function("HashMap: insert_cold_sync", |b| {
        b.iter_custom(|iters| {
            let hashmap: HashMap<u64, u64> = HashMap::default();
            let start = Instant::now();
            for i in 0..iters {
                assert!(hashmap.insert_sync(i, i).is_ok());
            }
            start.elapsed()
        })
    });
}

fn insert_warmed_up_async(c: &mut Criterion) {
    c.bench_function("HashMap: insert_warmed_up_async", |b| {
        b.to_async(FuturesExecutor).iter_custom(async |iters| {
            let hashmap: HashMap<u64, u64> = HashMap::with_capacity(iters as usize * 2);
            let start = Instant::now();
            for i in 0..iters {
                assert!(hashmap.insert_async(i, i).await.is_ok());
            }
            start.elapsed()
        })
    });
}

fn insert_warmed_up_sync(c: &mut Criterion) {
    c.bench_function("HashMap: insert_warmed_up_sync", |b| {
        b.iter_custom(|iters| {
            let hashmap: HashMap<u64, u64> = HashMap::with_capacity(iters as usize * 2);
            let start = Instant::now();
            for i in 0..iters {
                assert!(hashmap.insert_sync(i, i).is_ok());
            }
            start.elapsed()
        })
    });
}

fn read_async(c: &mut Criterion) {
    c.bench_function("HashMap: read_async", |b| {
        b.to_async(FuturesExecutor).iter_custom(async |iters| {
            let hashmap: HashMap<u64, u64> = HashMap::with_capacity(iters as usize * 2);
            for i in 0..iters {
                assert!(hashmap.insert_async(i, i).await.is_ok());
            }
            let start = Instant::now();
            for i in 0..iters {
                assert_eq!(hashmap.read_async(&i, |_, v| *v == i).await, Some(true));
            }
            start.elapsed()
        })
    });
}

fn read_sync(c: &mut Criterion) {
    c.bench_function("HashMap: read_sync", |b| {
        b.iter_custom(|iters| {
            let hashmap: HashMap<u64, u64> = HashMap::with_capacity(iters as usize * 2);
            for i in 0..iters {
                assert!(hashmap.insert_sync(i, i).is_ok());
            }
            let start = Instant::now();
            for i in 0..iters {
                assert_eq!(hashmap.read_sync(&i, |_, v| *v == i), Some(true));
            }
            start.elapsed()
        })
    });
}

fn insert_tail_latency(c: &mut Criterion) {
    c.bench_function("HashMap: insert, tail latency", move |b| {
        b.iter_custom(|iters| {
            let mut agg_max_latency = Duration::default();
            for _ in 0..iters {
                let hashmap: HashMap<u64, u64> = HashMap::default();
                let mut key = 0;
                let mut max_latency = Duration::default();
                (0..1048576).for_each(|_| {
                    key += 1;
                    let start = Instant::now();
                    assert!(hashmap.insert_sync(key, key).is_ok());
                    let elapsed = start.elapsed();
                    if elapsed > max_latency {
                        max_latency = elapsed;
                    }
                });
                agg_max_latency += max_latency;
            }
            agg_max_latency
        })
    });
}

fn throughput(c: &mut Criterion) {
    let capacity: usize = 16384;
    let num_ops: u64 = 65536;
    let num_threads: usize = 4;

    let hashmap = HashMap::with_capacity(capacity);
    for key in 0..capacity as u64 {
        assert!(hashmap.insert_sync(key, key).is_ok());
    }

    // 550us: target = 400us.
    let mut read = c.benchmark_group("read_throughput");
    read.throughput(Throughput::Elements(num_ops));
    read.bench_function("HashMap: read, throughput", |b| {
        b.iter(|| {
            thread::scope(|scope| {
                let mut threads = Vec::with_capacity(num_threads);
                let hashmap = &hashmap;
                for thread_id in 0..num_threads {
                    threads.push(scope.spawn(move || {
                        let mut sum: u64 = 0;
                        for k in (0..num_ops / num_threads as u64)
                            .map(|k| k * num_threads as u64 + thread_id as u64)
                        {
                            if let Some(v) = hashmap.read_sync(&k, |_, v| *v) {
                                sum = sum.wrapping_add(v);
                            }
                        }
                        black_box(sum);
                    }));
                }
                for thread in threads {
                    assert!(thread.join().is_ok());
                }
            })
        });
    });
    read.finish();

    // 750us: target = 400us.
    let mut read_write = c.benchmark_group("read_write_throughput");
    read_write.throughput(Throughput::Elements(num_ops));
    read_write.bench_function("HashMap: read 0.9 write 0.1, throughput", |b| {
        b.iter(|| {
            let mut sum: u64 = 0;
            for k in 0..num_ops {
                if k % 10 == 0 {
                    if let Some(v) = hashmap.update_sync(&k, |_, v| {
                        *v = k.wrapping_add(k);
                        *v
                    }) {
                        sum = sum.wrapping_add(v);
                    }
                } else if let Some(v) = hashmap.read_sync(&k, |_, v| *v) {
                    sum = sum.wrapping_add(v);
                }
            }
            black_box(sum)
        });
    });
    read_write.finish();

    // 350us: target = 200us.
    let mut insert = c.benchmark_group("insert_throughput");
    insert.throughput(Throughput::Elements(capacity as u64));
    insert.bench_function("HashMap: insert, throughput", |b| {
        b.iter(|| {
            let hashmap = HashMap::with_capacity(capacity);
            for k in 0..capacity as u64 {
                assert!(hashmap.insert_sync(k, k).is_ok());
            }
            black_box(hashmap.len())
        });
    });
    insert.finish();
}

criterion_group!(
    hash_map,
    insert_remove_single_async,
    insert_remove_single_sync,
    insert_cold_async,
    insert_cold_sync,
    insert_warmed_up_async,
    insert_warmed_up_sync,
    insert_tail_latency,
    read_async,
    read_sync,
    throughput
);
criterion_main!(hash_map);
