// benches/memory_benchmarks.rs
//
// Criterion benchmarks that call the MemoryStorage layer directly, bypassing HTTP.
// These measure the pure in-memory cost of each operation — no disk I/O, no fsync.
//
// Run with:
//   cargo bench --bench memory_benchmarks
//
// Compare with storage_benchmarks to see the RocksDB / fsync overhead.
// Compare with http_memory_benchmarks to see the web-layer overhead on top of
// in-memory storage.

use chrono::Utc;
use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
use qrusty::{
    api::StorageApi,
    memory_storage::MemoryStorage,
    message::{Message, Priority, PriorityOrdering, QueueConfig},
};
use std::sync::atomic::{AtomicU64, Ordering};

const DEPTHS: &[usize] = &[0, 100, 1_000, 10_000];
const BATCH_SIZES: &[usize] = &[1, 10, 50, 100];
const BATCH_DEPTH: usize = 1_000;

const QUEUE_ALLOW_DUP: &str = "mem_bench_allow_dup";
const QUEUE_NO_DUP: &str = "mem_bench_no_dup";
const CONSUMER: &str = "bench-consumer";

static COUNTER: AtomicU64 = AtomicU64::new(0);

fn next_id() -> u64 {
    COUNTER.fetch_add(1, Ordering::Relaxed)
}

fn make_message(queue: &str, payload: &str, priority: u64) -> Message {
    Message {
        id: format!("mem-bench-{}", next_id()),
        queue: queue.to_string(),
        priority: Priority::Numeric(priority),
        payload: payload.to_string(),
        created_at: Utc::now(),
        locked_until: None,
        locked_by: None,
        retry_count: 0,
        max_retries: 10,
        payload_ref: None,
        payload_hash: None,
    }
}

/// Creates a MemoryStorage instance and pre-populates `depth` background messages.
fn setup_memory(
    rt: &tokio::runtime::Runtime,
    queue: &str,
    allow_duplicates: bool,
    depth: usize,
) -> MemoryStorage {
    let storage = MemoryStorage::new();

    rt.block_on(async {
        storage
            .create_queue(
                queue,
                QueueConfig {
                    ordering: PriorityOrdering::MaxFirst,
                    allow_duplicates,
                    ..Default::default()
                },
            )
            .await;

        for i in 0..depth {
            let msg = make_message(queue, &format!("background-{}", i), 100);
            storage.push(msg).await.expect("pre-populate push");
        }
    });

    storage
}

// ---------------------------------------------------------------------------
// Scenario 1: Enqueue — allow_duplicates = true
// ---------------------------------------------------------------------------
fn bench_push_allow_duplicates(c: &mut Criterion) {
    let rt = tokio::runtime::Runtime::new().unwrap();
    let mut group = c.benchmark_group("memory/push/allow_duplicates");
    group.sample_size(200);

    for &depth in DEPTHS {
        let storage = setup_memory(&rt, QUEUE_ALLOW_DUP, true, depth);

        group.bench_with_input(
            BenchmarkId::new("queue_depth", depth),
            &depth,
            |b, _depth| {
                b.iter(|| {
                    let n = next_id();
                    let msg = make_message(QUEUE_ALLOW_DUP, &format!("payload-{}", n), 50);
                    rt.block_on(storage.push(msg)).expect("push")
                });
            },
        );
    }

    group.finish();
}

// ---------------------------------------------------------------------------
// Scenario 2: Enqueue — allow_duplicates = false
//
// The MemoryStorage duplicate check scans the HashSet (O(1) per lookup) rather
// than a linear prefix scan. At high queue depths this is markedly faster than
// the RocksDB equivalent.
// ---------------------------------------------------------------------------
fn bench_push_no_duplicates(c: &mut Criterion) {
    let rt = tokio::runtime::Runtime::new().unwrap();
    let mut group = c.benchmark_group("memory/push/no_duplicates");
    group.sample_size(200);

    for &depth in DEPTHS {
        let storage = setup_memory(&rt, QUEUE_NO_DUP, false, depth);

        group.bench_with_input(
            BenchmarkId::new("queue_depth", depth),
            &depth,
            |b, _depth| {
                b.iter(|| {
                    let n = next_id();
                    let msg = make_message(QUEUE_NO_DUP, &format!("unique-{}", n), 50);
                    rt.block_on(storage.push(msg)).expect("push")
                });
            },
        );
    }

    group.finish();
}

// ---------------------------------------------------------------------------
// Scenario 3: ACK
// ---------------------------------------------------------------------------
fn bench_ack(c: &mut Criterion) {
    let rt = tokio::runtime::Runtime::new().unwrap();
    let mut group = c.benchmark_group("memory/ack");
    group.sample_size(200);

    for &depth in DEPTHS {
        let storage = setup_memory(&rt, QUEUE_ALLOW_DUP, true, depth);

        group.bench_with_input(
            BenchmarkId::new("queue_depth", depth),
            &depth,
            |b, _depth| {
                b.iter_batched(
                    || {
                        rt.block_on(async {
                            let n = next_id();
                            let msg =
                                make_message(QUEUE_ALLOW_DUP, &format!("ack-target-{}", n), 1);
                            storage.push(msg).await.expect("setup push");
                            storage
                                .pop(QUEUE_ALLOW_DUP, CONSUMER, 3600)
                                .await
                                .expect("setup pop")
                                .expect("no message to pop")
                        })
                    },
                    |locked_msg| {
                        rt.block_on(storage.ack(QUEUE_ALLOW_DUP, &locked_msg.id, CONSUMER))
                            .expect("ack")
                    },
                    BatchSize::SmallInput,
                );
            },
        );
    }

    group.finish();
}

// ---------------------------------------------------------------------------
// Scenario 4a: NACK — allow_duplicates = true
// ---------------------------------------------------------------------------
fn bench_nack_allow_duplicates(c: &mut Criterion) {
    let rt = tokio::runtime::Runtime::new().unwrap();
    let mut group = c.benchmark_group("memory/nack/allow_duplicates");
    group.sample_size(200);

    for &depth in DEPTHS {
        let storage = setup_memory(&rt, QUEUE_ALLOW_DUP, true, depth);

        group.bench_with_input(
            BenchmarkId::new("queue_depth", depth),
            &depth,
            |b, _depth| {
                b.iter_batched(
                    || {
                        rt.block_on(async {
                            let n = next_id();
                            let msg =
                                make_message(QUEUE_ALLOW_DUP, &format!("nack-target-{}", n), 1);
                            storage.push(msg).await.expect("setup push");
                            storage
                                .pop(QUEUE_ALLOW_DUP, CONSUMER, 3600)
                                .await
                                .expect("setup pop")
                                .expect("no message to pop")
                        })
                    },
                    |locked_msg| {
                        rt.block_on(storage.nack(QUEUE_ALLOW_DUP, &locked_msg.id, CONSUMER))
                            .expect("nack")
                    },
                    BatchSize::SmallInput,
                );
            },
        );
    }

    group.finish();
}

// ---------------------------------------------------------------------------
// Scenario 4b: NACK — allow_duplicates = false
// ---------------------------------------------------------------------------
fn bench_nack_no_duplicates(c: &mut Criterion) {
    let rt = tokio::runtime::Runtime::new().unwrap();
    let mut group = c.benchmark_group("memory/nack/no_duplicates");
    group.sample_size(200);

    for &depth in DEPTHS {
        let storage = setup_memory(&rt, QUEUE_NO_DUP, false, depth);

        group.bench_with_input(
            BenchmarkId::new("queue_depth", depth),
            &depth,
            |b, _depth| {
                b.iter_batched(
                    || {
                        rt.block_on(async {
                            let n = next_id();
                            let msg = make_message(QUEUE_NO_DUP, &format!("nack-unique-{}", n), 1);
                            storage.push(msg).await.expect("setup push");
                            storage
                                .pop(QUEUE_NO_DUP, CONSUMER, 3600)
                                .await
                                .expect("setup pop")
                                .expect("no message to pop")
                        })
                    },
                    |locked_msg| {
                        rt.block_on(storage.nack(QUEUE_NO_DUP, &locked_msg.id, CONSUMER))
                            .expect("nack")
                    },
                    BatchSize::SmallInput,
                );
            },
        );
    }

    group.finish();
}

// ---------------------------------------------------------------------------
// Scenario 5: Batch ACK
// ---------------------------------------------------------------------------
fn bench_batch_ack(c: &mut Criterion) {
    let rt = tokio::runtime::Runtime::new().unwrap();
    let mut group = c.benchmark_group("memory/batch_ack");
    group.sample_size(50);

    for &batch_size in BATCH_SIZES {
        let storage = setup_memory(&rt, QUEUE_ALLOW_DUP, true, BATCH_DEPTH);

        group.bench_with_input(
            BenchmarkId::new("batch_size", batch_size),
            &batch_size,
            |b, &bs| {
                b.iter_batched(
                    || {
                        rt.block_on(async {
                            let mut ids = Vec::with_capacity(bs);
                            for _ in 0..bs {
                                let n = next_id();
                                let msg =
                                    make_message(QUEUE_ALLOW_DUP, &format!("batch-ack-{}", n), 1);
                                storage.push(msg).await.expect("setup push");
                                let locked = storage
                                    .pop(QUEUE_ALLOW_DUP, CONSUMER, 3600)
                                    .await
                                    .expect("setup pop")
                                    .expect("no message");
                                ids.push(locked.id);
                            }
                            ids
                        })
                    },
                    |ids| {
                        rt.block_on(storage.batch_ack(QUEUE_ALLOW_DUP, CONSUMER, &ids))
                            .expect("batch_ack")
                    },
                    BatchSize::SmallInput,
                );
            },
        );
    }

    group.finish();
}

// ---------------------------------------------------------------------------
// Scenario 6a: Batch NACK — allow_duplicates = true
// ---------------------------------------------------------------------------
fn bench_batch_nack_allow_duplicates(c: &mut Criterion) {
    let rt = tokio::runtime::Runtime::new().unwrap();
    let mut group = c.benchmark_group("memory/batch_nack/allow_duplicates");
    group.sample_size(50);

    for &batch_size in BATCH_SIZES {
        let storage = setup_memory(&rt, QUEUE_ALLOW_DUP, true, BATCH_DEPTH);

        group.bench_with_input(
            BenchmarkId::new("batch_size", batch_size),
            &batch_size,
            |b, &bs| {
                b.iter_batched(
                    || {
                        rt.block_on(async {
                            let mut ids = Vec::with_capacity(bs);
                            for _ in 0..bs {
                                let n = next_id();
                                let msg = make_message(
                                    QUEUE_ALLOW_DUP,
                                    &format!("batch-nack-dup-{}", n),
                                    1,
                                );
                                storage.push(msg).await.expect("setup push");
                                let locked = storage
                                    .pop(QUEUE_ALLOW_DUP, CONSUMER, 3600)
                                    .await
                                    .expect("setup pop")
                                    .expect("no message");
                                ids.push(locked.id);
                            }
                            ids
                        })
                    },
                    |ids| {
                        rt.block_on(storage.batch_nack(QUEUE_ALLOW_DUP, CONSUMER, &ids))
                            .expect("batch_nack")
                    },
                    BatchSize::SmallInput,
                );
            },
        );
    }

    group.finish();
}

// ---------------------------------------------------------------------------
// Scenario 6b: Batch NACK — allow_duplicates = false
// ---------------------------------------------------------------------------
fn bench_batch_nack_no_duplicates(c: &mut Criterion) {
    let rt = tokio::runtime::Runtime::new().unwrap();
    let mut group = c.benchmark_group("memory/batch_nack/no_duplicates");
    group.sample_size(50);

    for &batch_size in BATCH_SIZES {
        let storage = setup_memory(&rt, QUEUE_NO_DUP, false, BATCH_DEPTH);

        group.bench_with_input(
            BenchmarkId::new("batch_size", batch_size),
            &batch_size,
            |b, &bs| {
                b.iter_batched(
                    || {
                        rt.block_on(async {
                            let mut ids = Vec::with_capacity(bs);
                            for _ in 0..bs {
                                let n = next_id();
                                let msg = make_message(
                                    QUEUE_NO_DUP,
                                    &format!("batch-nack-nodup-{}", n),
                                    1,
                                );
                                storage.push(msg).await.expect("setup push");
                                let locked = storage
                                    .pop(QUEUE_NO_DUP, CONSUMER, 3600)
                                    .await
                                    .expect("setup pop")
                                    .expect("no message");
                                ids.push(locked.id);
                            }
                            ids
                        })
                    },
                    |ids| {
                        rt.block_on(storage.batch_nack(QUEUE_NO_DUP, CONSUMER, &ids))
                            .expect("batch_nack")
                    },
                    BatchSize::SmallInput,
                );
            },
        );
    }

    group.finish();
}

criterion_group!(
    benches,
    bench_push_allow_duplicates,
    bench_push_no_duplicates,
    bench_ack,
    bench_nack_allow_duplicates,
    bench_nack_no_duplicates,
    bench_batch_ack,
    bench_batch_nack_allow_duplicates,
    bench_batch_nack_no_duplicates,
);
criterion_main!(benches);
