use crate::{ITEM_SIZE, ITEMS_PER_BLOB, append_fixed_random_data, get_variable_journal};
use commonware_runtime::{
    ReadOptions, Runner as _, Supervisor as _,
    benchmarks::{context, tokio},
    tokio::{Config, Context, Runner},
};
use commonware_storage::journal::contiguous::{Contiguous as _, variable::Journal};
use commonware_utils::{NZUsize, sequence::FixedBytes};
use criterion::{Criterion, criterion_group};
use futures::{StreamExt, pin_mut};
use std::{
    hint::black_box,
    time::{Duration, Instant},
};

/// Partition name to use in the journal config.
const PARTITION: &str = "variable-test-partition";

/// Replay all items in the given `journal`.
async fn bench_run(
    journal: Journal<Context, FixedBytes<ITEM_SIZE>>,
    buffer: usize,
    read_options: ReadOptions,
) -> Journal<Context, FixedBytes<ITEM_SIZE>> {
    let (journal, reader) = journal.snapshot().await.unwrap();
    let stream = reader
        .replay(0, NZUsize!(buffer), read_options)
        .await
        .expect("failed to replay journal");
    pin_mut!(stream);
    while let Some(result) = stream.next().await {
        match result {
            Ok(item) => {
                black_box(item);
            }
            Err(err) => panic!("Failed to read item: {err}"),
        }
    }
    journal
}

/// Benchmark the replaying of items from a variable journal containing exactly that
/// number of items.
fn bench_variable_replay(c: &mut Criterion) {
    for items in [1_000, 10_000, 100_000, 500_000] {
        let cfg = Config::default();
        let mut initialized = false;
        let runner = tokio::Runner::new(cfg.clone());

        // Run DONT_CACHE first because it may not prune pages left resident by
        // cached replays.
        for (read_options, label) in [
            (ReadOptions::DONT_CACHE, "dont_cache"),
            (ReadOptions::default(), "cache"),
        ] {
            for buffer in [16_384, 65_536, 1_048_576] {
                c.bench_function(
                    &format!(
                        "{}/items={} buffer={} size={} read_options={}",
                        module_path!(),
                        items,
                        buffer,
                        ITEM_SIZE,
                        label
                    ),
                    |b| {
                        // Setup: populate journal (once, on first sample).
                        if !initialized {
                            Runner::new(cfg.clone()).start(|ctx| async move {
                                let j = get_variable_journal(ctx, PARTITION, ITEMS_PER_BLOB).await;
                                append_fixed_random_data::<_, ITEM_SIZE>(j, items).await;
                            });
                            initialized = true;
                        }

                        // Benchmark: measure replay time.
                        b.to_async(&runner).iter_custom(|iters| async move {
                            let ctx = context::get::<commonware_runtime::tokio::Context>();
                            let mut j = get_variable_journal(
                                ctx.child("storage"),
                                PARTITION,
                                ITEMS_PER_BLOB,
                            )
                            .await;
                            let mut duration = Duration::ZERO;
                            for _ in 0..iters {
                                let start = Instant::now();
                                j = bench_run(j, buffer, read_options).await;
                                duration += start.elapsed();
                            }
                            duration
                        });
                    },
                );
            }
        }

        // Cleanup: destroy journal.
        if initialized {
            Runner::new(cfg).start(|context| async move {
                let j = get_variable_journal::<ITEM_SIZE>(context, PARTITION, ITEMS_PER_BLOB).await;
                j.destroy().await.unwrap();
            });
        }
    }
}

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