use std::hint::black_box;
use std::iter;
use std::sync::atomic::AtomicBool;

use common::universal_io::{MmapFile, MmapFs};
use criterion::{Criterion, criterion_group, criterion_main};
use rand::rngs::StdRng;
use rand::{RngExt, SeedableRng};
use segment::common::flags::dynamic_stored_flags::DynamicStoredFlags;
use segment::common::operation_error::check_process_stopped;
use tempfile::tempdir;

const FLAG_COUNT: usize = 50_000_000;

fn dynamic_mmap_flag_count(c: &mut Criterion) {
    let mut rng = StdRng::seed_from_u64(42);
    let dir = tempdir().unwrap();
    let random_flags: Vec<bool> = iter::repeat_with(|| rng.random())
        .take(FLAG_COUNT)
        .collect();
    let stopped = AtomicBool::new(false);

    // Build dynamic mmap flags with random deletions
    let mut dynamic_flags =
        DynamicStoredFlags::<MmapFile>::open(&MmapFs, dir.path(), false).unwrap();
    dynamic_flags.set_len(&MmapFs, FLAG_COUNT).unwrap();
    dynamic_flags
        .set_ascending_bits(
            random_flags
                .iter()
                .enumerate()
                .filter(|(_, flag)| **flag)
                .map(|(i, _)| (i as u64, true)),
        )
        .unwrap();

    dynamic_flags.flusher()().unwrap();
    let real_count = random_flags.iter().filter(|&&flag| flag).count();

    let mut group = c.benchmark_group("dynamic-mmap-flag-count");

    group.bench_function("manual-count-loop-stoppable", |b| {
        b.iter(|| {
            let mut count = 0;
            for i in 0..FLAG_COUNT {
                if dynamic_flags.get(i).unwrap() {
                    count += 1;
                }
                check_process_stopped(&stopped).unwrap();
            }
            assert_eq!(count, real_count);
            black_box(count)
        });
    });

    group.bench_function("manual-count-loop", |b| {
        b.iter(|| {
            let mut count = 0;
            for i in 0..FLAG_COUNT {
                if dynamic_flags.get(i).unwrap() {
                    count += 1;
                }
            }
            assert_eq!(count, real_count);
            black_box(count)
        });
    });

    group.bench_function("count-ones", |b| {
        b.iter(|| {
            let count = dynamic_flags.count_flags().unwrap();
            assert_eq!(count, real_count);
            black_box(count)
        });
    });
}

criterion_group! {
    name = benches;
    config = Criterion::default();
    targets = dynamic_mmap_flag_count
}

criterion_main!(benches);
