use std::{borrow::BorrowMut, collections::BTreeSet};

use criterion::{black_box, criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
use heckcheck::prelude::*;

struct Slab {
    index: BTreeSet<usize>,
    slab: slab::Slab<usize>,
}

impl Slab {
    fn new() -> Self {
        Self {
            index: BTreeSet::new(),
            slab: slab::Slab::new(),
        }
    }

    fn with_capacity(cap: usize) -> Self {
        Self {
            index: BTreeSet::new(),
            slab: slab::Slab::with_capacity(cap),
        }
    }

    fn insert(&mut self, value: usize) {
        self.index.insert(value);
        self.slab.insert(value);
    }

    fn reserve(&mut self, value: usize) {
        self.slab.reserve(value);
    }

    fn contains(&self, value: usize) -> bool {
        self.index.contains(&value)
    }

    fn get(&self, value: usize) -> Option<&usize> {
        self.index.get(&value)
    }

    fn remove(&mut self, value: usize) {
        self.index.remove(&value);
        self.slab.try_remove(value);
    }

    fn clear(&mut self) {
        self.index.clear();
        self.slab.clear();
    }

    fn iter(&mut self) -> impl Iterator<Item = (usize, usize)> + '_ {
        self.index
            .iter()
            .map(|key| (*key, *self.slab.borrow_mut().get(*key).unwrap()))
    }
}

fn insert(c: &mut Criterion) {
    let mut group = c.benchmark_group("insert (with index)");
    for i in [10, 100, 1_000, 10_000].iter() {
        group.bench_with_input(BenchmarkId::new("beton", i), i, |b, i| {
            let setup = || beton::Slab::with_capacity(*i);
            let routine = |mut slab: beton::Slab<usize>| {
                black_box({
                    for n in 0..*i {
                        slab.insert(n);
                    }
                });
            };
            b.iter_batched(setup, routine, BatchSize::SmallInput)
        });

        group.bench_with_input(BenchmarkId::new("slab", i), i, |b, i| {
            let setup = || Slab::with_capacity(*i);
            let routine = |mut slab: Slab| {
                black_box({
                    for n in 0..*i {
                        slab.insert(n);
                    }
                });
            };
            b.iter_batched(setup, routine, BatchSize::SmallInput)
        });
    }
    group.finish();
}

fn mutate(c: &mut Criterion) {
    let mut group = c.benchmark_group("mutate (with index)");
    let seed = fastrand::u64(1024..2048);

    /// A single operation we can apply
    #[derive(Arbitrary, Debug)]
    enum Operation {
        Insert(usize),
        Fetch(usize),
        Remove(usize),
        Contains(usize),
        Clear,
        Reserve(usize),
    }

    group.bench_function(BenchmarkId::new("beton", "main"), |b| {
        let setup = || beton::Slab::new();
        let routine = |mut slab: beton::Slab<usize>| {
            black_box({
                let mut checker = heckcheck::HeckCheck::from_seed(seed);
                checker.check(|operations: Vec<Operation>| {
                    for operation in operations {
                        match operation {
                            Operation::Insert(value) => {
                                let _key = slab.insert(value);
                            }
                            Operation::Fetch(index) => {
                                let _output = slab.get(index.into());
                            }
                            Operation::Remove(index) => {
                                let _output = slab.remove(index.into());
                            }
                            Operation::Contains(index) => {
                                let _contains = slab.contains_key(index.into());
                            }
                            Operation::Clear => {
                                slab.clear();
                            }
                            Operation::Reserve(capacity) => {
                                // NOTE: very big allocations make this slow, so we ensure
                                // they're always a little smaller
                                //
                                // FIXME: in the proper impl of BitTree, we should be able
                                // to a `memset(3)` call.
                                let capacity = capacity % 1024;
                                slab.reserve(capacity);
                            }
                        }
                    }
                    Ok(())
                });
            });
        };
        b.iter_batched(setup, routine, BatchSize::SmallInput)
    });

    group.bench_function(BenchmarkId::new("slab", "main"), |b| {
        let setup = || Slab::new();
        let routine = |mut slab: Slab| {
            black_box({
                let mut checker = heckcheck::HeckCheck::from_seed(seed);
                checker.check(|operations: Vec<Operation>| {
                    for operation in operations {
                        match operation {
                            Operation::Insert(value) => {
                                let _key = slab.insert(value);
                            }
                            Operation::Fetch(index) => {
                                let _output = slab.get(index);
                            }
                            Operation::Remove(index) => {
                                let _output = slab.remove(index);
                            }
                            Operation::Contains(index) => {
                                let _contains = slab.contains(index);
                            }
                            Operation::Clear => {
                                slab.clear();
                            }
                            Operation::Reserve(capacity) => {
                                // NOTE: very big allocations make this slow, so we ensure
                                // they're always a little smaller
                                let capacity = capacity % 1024;
                                slab.reserve(capacity);
                            }
                        }
                    }
                    Ok(())
                });
            });
        };
        b.iter_batched(setup, routine, BatchSize::SmallInput)
    });
    group.finish();
}

fn iterate(c: &mut Criterion) {
    let mut group = c.benchmark_group("iterate (with index)");
    for i in [10, 100, 1_000, 10_000].iter() {
        group.bench_with_input(BenchmarkId::new("beton", i), i, |b, i| {
            let setup = || {
                let mut slab = beton::Slab::with_capacity(*i);
                let mut total = 0;
                for n in 0..*i {
                    total += n;
                    slab.insert(n);
                }
                (slab, total)
            };
            let routine = |(slab, total): (beton::Slab<usize>, usize)| {
                black_box({
                    let mut sum = 0;
                    for (_, n) in slab {
                        sum += n;
                    }
                    assert_eq!(sum, total);
                });
            };
            b.iter_batched(setup, routine, BatchSize::SmallInput)
        });

        group.bench_with_input(BenchmarkId::new("slab", i), i, |b, i| {
            let setup = || {
                let mut slab = Slab::with_capacity(*i);
                let mut total = 0;
                for n in 0..*i {
                    total += n;
                    slab.insert(n);
                }
                (slab, total)
            };
            let routine = |(mut slab, total): (Slab, usize)| {
                black_box({
                    let mut sum = 0;
                    for (_, n) in slab.iter() {
                        sum += n;
                    }
                    assert_eq!(sum, total);
                });
            };
            b.iter_batched(setup, routine, BatchSize::SmallInput)
        });
    }
    group.finish();
}

criterion_group!(operations_with_index, insert, mutate, iterate);
criterion_main!(operations_with_index);