use super::utils::{append_random, init, Ordinal, ITEMS_PER_BLOB};
use commonware_runtime::{
    benchmarks::{context, tokio},
    tokio::Config,
    Runner,
};
use commonware_storage::utils::bits_for_indices;
use commonware_utils::{TestRng, NZU64};
use criterion::{criterion_group, Criterion};
use futures::{stream::FuturesUnordered, StreamExt};
use rand::RngExt as _;
use std::{hint::black_box, time::Instant};

/// Items pre-loaded into the store.
const ITEMS: u64 = 250_000;

/// Select random indices for benchmarking.
pub fn select_indices(count: usize, items: u64) -> Vec<u64> {
    let mut rng = TestRng::new(42);
    let mut selected_indices = Vec::with_capacity(count);
    for _ in 0..count {
        selected_indices.push(rng.random_range(0..items));
    }
    selected_indices
}

/// Read indices serially from an ordinal store.
pub async fn read_serial_indices(store: &Ordinal, indices: &[u64]) {
    for idx in indices {
        black_box(store.get(*idx).await.unwrap().unwrap());
    }
}

/// Read indices concurrently from an ordinal store.
pub async fn read_concurrent_indices(store: &Ordinal, indices: &[u64]) {
    let mut futures = FuturesUnordered::new();
    for idx in indices {
        futures.push(store.get(*idx));
    }
    while let Some(result) = futures.next().await {
        black_box(result.unwrap().unwrap());
    }
}

fn bench_get(c: &mut Criterion) {
    // Create a config we can use across all benchmarks (with a fixed `storage_directory`).
    let cfg = Config::default();

    // Create a shared on-disk store once so later setup is fast.
    let builder = commonware_runtime::tokio::Runner::new(cfg.clone());
    let bits = builder.start(|ctx| async move {
        let mut store = init(ctx, None).await;
        let indices = append_random(&mut store, ITEMS).await;
        store.sync().await.unwrap();
        bits_for_indices(NZU64!(ITEMS_PER_BLOB), indices)
    });

    // Run the benchmarks.
    let runner = tokio::Runner::new(cfg.clone());
    for mode in ["serial", "concurrent"] {
        for reads in [1_000, 10_000, 50_000] {
            let label = format!("{}/mode={} reads={}", module_path!(), mode, reads);
            c.bench_function(&label, |b| {
                b.to_async(&runner).iter_custom({
                    let bits = bits.clone();
                    move |iters| {
                        let bits = bits.clone();
                        async move {
                            let ctx = context::get::<commonware_runtime::tokio::Context>();
                            let bits = bits
                                .iter()
                                .map(|(section, bitmap)| (*section, bitmap))
                                .collect();
                            let store = init(ctx, Some(bits)).await;
                            let selected_indices = select_indices(reads, ITEMS);
                            let start = Instant::now();
                            for _ in 0..iters {
                                match mode {
                                    "serial" => {
                                        read_serial_indices(&store, &selected_indices).await
                                    }
                                    "concurrent" => {
                                        read_concurrent_indices(&store, &selected_indices).await
                                    }
                                    _ => unreachable!(),
                                }
                            }
                            start.elapsed()
                        }
                    }
                });
            });
        }
    }

    // Clean up shared artifacts.
    let cleaner = commonware_runtime::tokio::Runner::new(cfg);
    cleaner.start(|ctx| async move {
        let store = init(ctx, None).await;
        store.destroy().await.unwrap();
    });
}

criterion_group! {
    name = benches;
    config = Criterion::default().sample_size(10);
    targets = bench_get
}