use bplus_index::BPlusIndex;
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
use rand::prelude::*;
use std::{collections::BTreeMap, hint::black_box};

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

fn shuffled_keys(n: usize) -> Vec<i64> {
    let mut rng = StdRng::seed_from_u64(42);
    let mut keys: Vec<i64> = (0..n as i64).collect();
    keys.shuffle(&mut rng);
    keys
}

// ---------------------------------------------------------------------------
// Insert
// ---------------------------------------------------------------------------

fn bench_insert(c: &mut Criterion) {
    let mut group = c.benchmark_group("insert");
    for &n in &[1_000, 10_000, 100_000] {
        let keys = shuffled_keys(n);

        group.bench_with_input(BenchmarkId::new("BPlusIndex", n), &keys, |b, keys| {
            b.iter(|| {
                let mut idx = BPlusIndex::<i64, u64>::new();
                for (i, k) in keys.iter().enumerate() {
                    idx.insert(*k, i as u64);
                }
                black_box(&idx);
            });
        });

        group.bench_with_input(BenchmarkId::new("BTreeMap", n), &keys, |b, keys| {
            b.iter(|| {
                let mut map = BTreeMap::<i64, Vec<u64>>::new();
                for (i, k) in keys.iter().enumerate() {
                    map.entry(*k).or_default().push(i as u64);
                }
                black_box(&map);
            });
        });
    }
    group.finish();
}

// ---------------------------------------------------------------------------
// Lookup (point query)
// ---------------------------------------------------------------------------

fn bench_lookup(c: &mut Criterion) {
    let mut group = c.benchmark_group("lookup");
    for &n in &[1_000, 10_000, 100_000] {
        let keys = shuffled_keys(n);

        // Pre-populate BPlusIndex
        let mut idx = BPlusIndex::<i64, u64>::new();
        for (i, k) in keys.iter().enumerate() {
            idx.insert(*k, i as u64);
        }

        // Pre-populate BTreeMap
        let mut map = BTreeMap::<i64, Vec<u64>>::new();
        for (i, k) in keys.iter().enumerate() {
            map.entry(*k).or_default().push(i as u64);
        }

        // Query keys in random order
        let query_keys = {
            let mut qk = keys.clone();
            qk.shuffle(&mut StdRng::seed_from_u64(99));
            qk
        };

        group.bench_with_input(BenchmarkId::new("BPlusIndex", n), &query_keys, |b, qk| {
            b.iter(|| {
                for k in qk {
                    black_box(idx.lookup_unique(k));
                }
            });
        });

        group.bench_with_input(BenchmarkId::new("BTreeMap", n), &query_keys, |b, qk| {
            b.iter(|| {
                for k in qk {
                    black_box(map.get(k));
                }
            });
        });
    }
    group.finish();
}

// ---------------------------------------------------------------------------
// Range scan
// ---------------------------------------------------------------------------

fn bench_range(c: &mut Criterion) {
    let mut group = c.benchmark_group("range_scan");
    let n = 100_000usize;
    let keys = shuffled_keys(n);

    let mut idx = BPlusIndex::<i64, u64>::new();
    for (i, k) in keys.iter().enumerate() {
        idx.insert(*k, i as u64);
    }

    let mut map = BTreeMap::<i64, Vec<u64>>::new();
    for (i, k) in keys.iter().enumerate() {
        map.entry(*k).or_default().push(i as u64);
    }

    // Scan windows of ~1% of the keyspace
    let window = (n / 100) as i64;
    let starts: Vec<i64> = (0..100).map(|i| i * window).collect();

    group.bench_function(BenchmarkId::new("BPlusIndex", n), |b| {
        b.iter(|| {
            let mut total = 0usize;
            for &s in &starts {
                total += idx.range(s..s + window).count();
            }
            black_box(total);
        });
    });

    group.bench_function(BenchmarkId::new("BTreeMap", n), |b| {
        b.iter(|| {
            let mut total = 0usize;
            for &s in &starts {
                total += map
                    .range(s..(s + window))
                    .map(|(_, v)| v.len())
                    .sum::<usize>();
            }
            black_box(total);
        });
    });

    group.finish();
}

// ---------------------------------------------------------------------------
// Delete
// ---------------------------------------------------------------------------

fn bench_delete(c: &mut Criterion) {
    let mut group = c.benchmark_group("delete");
    for &n in &[1_000, 10_000, 100_000] {
        let keys = shuffled_keys(n);

        group.bench_with_input(BenchmarkId::new("BPlusIndex", n), &keys, |b, keys| {
            b.iter_batched(
                || {
                    let mut idx = BPlusIndex::<i64, u64>::new();
                    for (i, k) in keys.iter().enumerate() {
                        idx.insert(*k, i as u64);
                    }
                    idx
                },
                |mut idx| {
                    for (i, k) in keys.iter().enumerate() {
                        idx.delete(k, &(i as u64));
                    }
                    black_box(&idx);
                },
                criterion::BatchSize::SmallInput,
            );
        });

        group.bench_with_input(BenchmarkId::new("BTreeMap", n), &keys, |b, keys| {
            b.iter_batched(
                || {
                    let mut map = BTreeMap::<i64, Vec<u64>>::new();
                    for (i, k) in keys.iter().enumerate() {
                        map.entry(*k).or_default().push(i as u64);
                    }
                    map
                },
                |mut map| {
                    for k in keys {
                        map.remove(k);
                    }
                    black_box(&map);
                },
                criterion::BatchSize::SmallInput,
            );
        });
    }
    group.finish();
}

// ---------------------------------------------------------------------------
// Iteration (full scan)
// ---------------------------------------------------------------------------

fn bench_iter(c: &mut Criterion) {
    let mut group = c.benchmark_group("iter_full");
    let n = 100_000usize;
    let keys = shuffled_keys(n);

    let mut idx = BPlusIndex::<i64, u64>::new();
    for (i, k) in keys.iter().enumerate() {
        idx.insert(*k, i as u64);
    }

    let mut map = BTreeMap::<i64, u64>::new();
    for (i, k) in keys.iter().enumerate() {
        map.insert(*k, i as u64);
    }

    group.bench_function(BenchmarkId::new("BPlusIndex", n), |b| {
        b.iter(|| {
            let total: u64 = idx.iter().map(|(_, v)| *v).sum();
            black_box(total);
        });
    });

    group.bench_function(BenchmarkId::new("BTreeMap", n), |b| {
        b.iter(|| {
            let total: u64 = map.values().copied().sum();
            black_box(total);
        });
    });

    group.finish();
}

criterion_group!(
    benches,
    bench_insert,
    bench_lookup,
    bench_range,
    bench_delete,
    bench_iter
);
criterion_main!(benches);
