// benches/storage_benchmarks.rs
//
// Criterion benchmarks that call the Storage layer directly, bypassing HTTP.
// These measure the pure RocksDB cost of each operation.
//
// Run with:
//   cargo bench --bench storage_benchmarks
//
// Compare with http_benchmarks to isolate web connection overhead.

use chrono::Utc;
use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
use qrusty::{
    message::{Message, Priority, PriorityOrdering, QueueConfig},
    storage::Storage,
};
use std::sync::atomic::{AtomicU64, Ordering};
use tempfile::TempDir;

// Queue depths to benchmark at. Each creates a separate benchmark case.
// Higher depths expose O(n) scan costs for duplicate checking and ID searches.
const DEPTHS: &[usize] = &[0, 100, 1_000, 10_000];

// Batch sizes for batch ACK/NACK benchmarks.
// Reveals how per-operation fixed costs (fsync, RwLock) amortize across N messages.
const BATCH_SIZES: &[usize] = &[1, 10, 50, 100];

// Fixed queue depth for batch benchmarks.
// 1 000 background messages keeps scan costs realistic without excessive setup time.
const BATCH_DEPTH: usize = 1_000;

const QUEUE_ALLOW_DUP: &str = "bench_allow_dup";
const QUEUE_NO_DUP: &str = "bench_no_dup";
const CONSUMER: &str = "bench-consumer";

// Shared atomic counter for generating unique payloads and message IDs
// across benchmark iterations without needing mutable state in closures.
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!("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, // high so NACK benchmarks don't hit DLQ
        payload_ref: None,
        payload_hash: None,
    }
}

/// Creates a Storage instance and pre-populates `depth` background messages.
/// Returns (Storage, TempDir) — TempDir must stay alive for the benchmark duration.
fn setup_storage(queue: &str, allow_duplicates: bool, depth: usize) -> (Storage, TempDir) {
    let temp_dir = TempDir::new().expect("tempdir");
    let storage = Storage::new(temp_dir.path().to_str().unwrap()).expect("storage");

    storage.configure_queue(
        queue,
        QueueConfig {
            ordering: PriorityOrdering::MaxFirst,
            allow_duplicates,
            ..Default::default()
        },
    );

    // Pre-populate with background messages at priority 100.
    // These are never consumed; they exist only to drive realistic scan costs.
    for i in 0..depth {
        let msg = make_message(queue, &format!("background-{}", i), 100);
        storage.push(msg).expect("pre-populate push");
    }

    (storage, temp_dir)
}

// ---------------------------------------------------------------------------
// Scenario 1: Enqueue — allow_duplicates = true
//
// Each iteration pushes one new message with a unique payload at priority 50
// (below the background priority-100 messages, so it sorts after them in the
// prefix iterator — this does NOT affect push cost, which is O(log n) for the
// RocksDB write itself.  No duplicate scan is performed.)
// ---------------------------------------------------------------------------
fn bench_push_allow_duplicates(c: &mut Criterion) {
    let mut group = c.benchmark_group("storage/push/allow_duplicates");
    group.sample_size(200);

    for &depth in DEPTHS {
        let (storage, _dir) = setup_storage(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);
                    storage.push(msg).expect("push")
                });
            },
        );
    }

    group.finish();
}

// ---------------------------------------------------------------------------
// Scenario 2: Enqueue — allow_duplicates = false
//
// Each iteration pushes a message with a unique payload. Before storing,
// Storage::push() calls has_duplicate_payload(), which scans ALL unlocked
// messages in the queue prefix. Cost is O(n) in queue depth.
// ---------------------------------------------------------------------------
fn bench_push_no_duplicates(c: &mut Criterion) {
    let mut group = c.benchmark_group("storage/push/no_duplicates");
    group.sample_size(200);

    for &depth in DEPTHS {
        let (storage, _dir) = setup_storage(QUEUE_NO_DUP, false, depth);

        group.bench_with_input(
            BenchmarkId::new("queue_depth", depth),
            &depth,
            |b, _depth| {
                b.iter(|| {
                    // Each payload is unique, so the duplicate scan completes
                    // without early exit — worst-case scan cost.
                    let n = next_id();
                    let msg = make_message(QUEUE_NO_DUP, &format!("unique-{}", n), 50);
                    storage.push(msg).expect("push")
                });
            },
        );
    }

    group.finish();
}

// ---------------------------------------------------------------------------
// Scenario 3: ACK
//
// The setup phase (not measured) pushes a low-priority message and pops it
// so it is locked.  The measured routine acks it.
//
// ACK performs a linear prefix scan to find the message by ID, so cost grows
// with queue depth.  The test message is at priority 1 (below background
// priority-100 messages), placing it at the end of the scan — worst case.
// ---------------------------------------------------------------------------
fn bench_ack(c: &mut Criterion) {
    let mut group = c.benchmark_group("storage/ack");
    group.sample_size(200);

    for &depth in DEPTHS {
        let (storage, _dir) = setup_storage(QUEUE_ALLOW_DUP, true, depth);

        group.bench_with_input(
            BenchmarkId::new("queue_depth", depth),
            &depth,
            |b, _depth| {
                b.iter_batched(
                    || {
                        // Setup (not measured): push + pop a message at priority 1
                        // so it sits after all depth background messages in scan order.
                        let n = next_id();
                        let msg = make_message(QUEUE_ALLOW_DUP, &format!("ack-target-{}", n), 1);
                        storage.push(msg).expect("setup push");
                        storage
                            .pop(QUEUE_ALLOW_DUP, CONSUMER, 3600)
                            .expect("setup pop")
                            .expect("no message to pop")
                    },
                    |locked_msg| {
                        // Measured: ack the locked message
                        storage
                            .ack(QUEUE_ALLOW_DUP, &locked_msg.id, CONSUMER)
                            .expect("ack")
                    },
                    BatchSize::SmallInput,
                );
            },
        );
    }

    group.finish();
}

// ---------------------------------------------------------------------------
// Scenario 4a: NACK — allow_duplicates = true
//
// NACK with duplicates allowed: finds message by ID scan, then unlocks it
// (no duplicate check).  Cost: O(n) scan for the ID search.
// ---------------------------------------------------------------------------
fn bench_nack_allow_duplicates(c: &mut Criterion) {
    let mut group = c.benchmark_group("storage/nack/allow_duplicates");
    group.sample_size(200);

    for &depth in DEPTHS {
        let (storage, _dir) = setup_storage(QUEUE_ALLOW_DUP, true, depth);

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

    group.finish();
}

// ---------------------------------------------------------------------------
// Scenario 4b: NACK — allow_duplicates = false
//
// After the ID scan, NACK calls has_duplicate_payload() before re-queuing.
// Cost: O(n) scan (ID search) + O(n) scan (duplicate check) = O(2n).
// ---------------------------------------------------------------------------
fn bench_nack_no_duplicates(c: &mut Criterion) {
    let mut group = c.benchmark_group("storage/nack/no_duplicates");
    group.sample_size(200);

    for &depth in DEPTHS {
        let (storage, _dir) = setup_storage(QUEUE_NO_DUP, false, depth);

        group.bench_with_input(
            BenchmarkId::new("queue_depth", depth),
            &depth,
            |b, _depth| {
                b.iter_batched(
                    || {
                        // Each iteration uses a unique payload so the
                        // duplicate check completes the full scan (worst case).
                        let n = next_id();
                        let msg = make_message(QUEUE_NO_DUP, &format!("nack-unique-{}", n), 1);
                        storage.push(msg).expect("setup push");
                        storage
                            .pop(QUEUE_NO_DUP, CONSUMER, 3600)
                            .expect("setup pop")
                            .expect("no message to pop")
                    },
                    |locked_msg| {
                        storage
                            .nack(QUEUE_NO_DUP, &locked_msg.id, CONSUMER)
                            .expect("nack")
                    },
                    BatchSize::SmallInput,
                );
            },
        );
    }

    group.finish();
}

// ---------------------------------------------------------------------------
// Scenario 5: Batch ACK
//
// Setup (not measured): push `batch_size` low-priority messages and pop each
// one to lock it.  The measured routine calls batch_ack once for all IDs.
//
// Compared to `bench_ack` (one call per message), this shows the per-message
// cost saving from amortizing the single WriteBatch + RwLock write.
// ---------------------------------------------------------------------------
fn bench_batch_ack(c: &mut Criterion) {
    let mut group = c.benchmark_group("storage/batch_ack");
    // Fewer samples because setup clones become expensive at batch_size=100.
    group.sample_size(50);

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

        group.bench_with_input(
            BenchmarkId::new("batch_size", batch_size),
            &batch_size,
            |b, &bs| {
                b.iter_batched(
                    || {
                        // Setup (not measured): push + pop bs messages so they are locked.
                        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).expect("setup push");
                            let locked = storage
                                .pop(QUEUE_ALLOW_DUP, CONSUMER, 3600)
                                .expect("setup pop")
                                .expect("no message");
                            ids.push(locked.id);
                        }
                        ids
                    },
                    |ids| {
                        // Measured: one batch_ack call for all `bs` messages.
                        storage
                            .batch_ack(QUEUE_ALLOW_DUP, CONSUMER, &ids)
                            .expect("batch_ack")
                    },
                    BatchSize::SmallInput,
                );
            },
        );
    }

    group.finish();
}

// ---------------------------------------------------------------------------
// Scenario 6a: Batch NACK — allow_duplicates = true
//
// One prefix scan + one WriteBatch unlock for all `batch_size` messages.
// Cost is dominated by the single fsync, not the message count.
// ---------------------------------------------------------------------------
fn bench_batch_nack_allow_duplicates(c: &mut Criterion) {
    let mut group = c.benchmark_group("storage/batch_nack/allow_duplicates");
    group.sample_size(50);

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

        group.bench_with_input(
            BenchmarkId::new("batch_size", batch_size),
            &batch_size,
            |b, &bs| {
                b.iter_batched(
                    || {
                        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).expect("setup push");
                            let locked = storage
                                .pop(QUEUE_ALLOW_DUP, CONSUMER, 3600)
                                .expect("setup pop")
                                .expect("no message");
                            ids.push(locked.id);
                        }
                        ids
                    },
                    |ids| {
                        storage
                            .batch_nack(QUEUE_ALLOW_DUP, CONSUMER, &ids)
                            .expect("batch_nack")
                    },
                    BatchSize::SmallInput,
                );
            },
        );
    }

    group.finish();
}

// ---------------------------------------------------------------------------
// Scenario 6b: Batch NACK — allow_duplicates = false
//
// The single prefix scan simultaneously builds the unlocked-payloads set for
// duplicate detection, so there is no per-message has_duplicate_payload() scan.
// Compare to `bench_nack_no_duplicates` which pays O(2n) per message.
// ---------------------------------------------------------------------------
fn bench_batch_nack_no_duplicates(c: &mut Criterion) {
    let mut group = c.benchmark_group("storage/batch_nack/no_duplicates");
    group.sample_size(50);

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

        group.bench_with_input(
            BenchmarkId::new("batch_size", batch_size),
            &batch_size,
            |b, &bs| {
                b.iter_batched(
                    || {
                        let mut ids = Vec::with_capacity(bs);
                        for _ in 0..bs {
                            let n = next_id();
                            // Unique payloads so no drops — worst-case duplicate scan.
                            let msg =
                                make_message(QUEUE_NO_DUP, &format!("batch-nack-nodup-{}", n), 1);
                            storage.push(msg).expect("setup push");
                            let locked = storage
                                .pop(QUEUE_NO_DUP, CONSUMER, 3600)
                                .expect("setup pop")
                                .expect("no message");
                            ids.push(locked.id);
                        }
                        ids
                    },
                    |ids| {
                        storage
                            .batch_nack(QUEUE_NO_DUP, CONSUMER, &ids)
                            .expect("batch_nack")
                    },
                    BatchSize::SmallInput,
                );
            },
        );
    }

    group.finish();
}

/// Benchmarks pop latency from the hot tier vs RocksDB fallback (cold).
///
/// "hot" pops from a queue whose hot tier is seeded.
/// "cold" pops from a queue where the hot tier was drained, forcing a
/// RocksDB prefix scan.
fn bench_pop_hot_vs_cold(c: &mut Criterion) {
    let mut group = c.benchmark_group("pop_hot_vs_cold");

    for &depth in &[100, 1_000, 10_000] {
        // Hot: pop from a pre-seeded hot tier.
        group.bench_with_input(BenchmarkId::new("hot", depth), &depth, |b, &depth| {
            b.iter_batched(
                || {
                    let (storage, dir) = setup_storage(QUEUE_ALLOW_DUP, true, depth);
                    (storage, dir)
                },
                |(storage, _dir)| {
                    let _ = storage.pop(QUEUE_ALLOW_DUP, CONSUMER, 3600);
                },
                BatchSize::SmallInput,
            );
        });

        // Cold: drain the hot tier, then pop (forces RocksDB scan).
        group.bench_with_input(BenchmarkId::new("cold", depth), &depth, |b, &depth| {
            b.iter_batched(
                || {
                    let (storage, dir) = setup_storage(QUEUE_ALLOW_DUP, true, depth);
                    // Drain the hot tier.
                    storage.shrink_hot_tiers(0);
                    (storage, dir)
                },
                |(storage, _dir)| {
                    let _ = storage.pop(QUEUE_ALLOW_DUP, CONSUMER, 3600);
                },
                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,
    bench_pop_hot_vs_cold,
);
criterion_main!(benches);
