use std::sync::Arc;

use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
use tokio::runtime::Runtime;

use fafb_q::storage::{FsStorage, StorageBackend};
use fafb_q::types::{
    MessageEnvelope, MessageId, NamespaceId, PublishedMessage, TopicName, TopicOffset, now_millis,
};

fn make_published(id: u64, ns: &str, topic: &str, payload_size: usize) -> PublishedMessage {
    PublishedMessage {
        envelope: Arc::new(MessageEnvelope {
            id: MessageId(id),
            namespace: NamespaceId(ns.into()),
            topic: TopicName(topic.into()),
            payload: vec![0xAB; payload_size],
            timestamp: now_millis(),
            headers: vec![],
            ttl_ms: 0,
        }),
        offset: TopicOffset(id),
    }
}

fn make_published_with_ts(
    id: u64,
    ns: &str,
    topic: &str,
    payload_size: usize,
    timestamp: u64,
) -> PublishedMessage {
    PublishedMessage {
        envelope: Arc::new(MessageEnvelope {
            id: MessageId(id),
            namespace: NamespaceId(ns.into()),
            topic: TopicName(topic.into()),
            payload: vec![0xAB; payload_size],
            timestamp,
            headers: vec![],
            ttl_ms: 0,
        }),
        offset: TopicOffset(id),
    }
}

fn make_storage(dir: &std::path::Path) -> Arc<FsStorage> {
    let storage = Arc::new(FsStorage::new(dir));
    let rt = Runtime::new().unwrap();
    rt.block_on(storage.init()).unwrap();
    storage
}

fn append_single_topic(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();
    let mut group = c.benchmark_group("storage_append_single_topic");

    for payload_size in [64, 1024, 65536] {
        group.bench_with_input(
            BenchmarkId::new("payload", payload_size),
            &payload_size,
            |b, &size| {
                let dir = tempfile::tempdir().unwrap();
                let storage = make_storage(dir.path());
                let ns = NamespaceId("bench-ns".into());
                let topic = TopicName("bench-topic".into());
                let mut id = 0u64;

                b.iter(|| {
                    let msg = make_published(id, "bench-ns", "bench-topic", size);
                    id += 1;
                    rt.block_on(storage.append(&ns, &topic, &msg)).unwrap();
                });
            },
        );
    }

    group.finish();
}

fn append_cross_topic(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();
    let mut group = c.benchmark_group("storage_append_cross_topic");
    group.sample_size(10);

    for num_topics in [10, 100] {
        group.bench_with_input(
            BenchmarkId::new("topics", num_topics),
            &num_topics,
            |b, &n| {
                let dir = tempfile::tempdir().unwrap();
                let storage = make_storage(dir.path());

                b.iter(|| {
                    rt.block_on(async {
                        let mut handles = Vec::with_capacity(n);
                        for topic_id in 0..n {
                            let storage = Arc::clone(&storage);
                            let ns = NamespaceId("bench-ns".into());
                            let topic = TopicName(format!("topic-{topic_id}"));
                            handles.push(tokio::spawn(async move {
                                for i in 0..100u64 {
                                    let msg = make_published(
                                        i,
                                        "bench-ns",
                                        &format!("topic-{topic_id}"),
                                        256,
                                    );
                                    storage.append(&ns, &topic, &msg).await.unwrap();
                                }
                            }));
                        }
                        for h in handles {
                            h.await.unwrap();
                        }
                    });
                });
            },
        );
    }

    group.finish();
}

fn replay_throughput(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();
    let mut group = c.benchmark_group("storage_replay");
    group.sample_size(10);

    for num_messages in [1_000, 10_000, 100_000] {
        group.bench_with_input(
            BenchmarkId::new("messages", num_messages),
            &num_messages,
            |b, &n| {
                let dir = tempfile::tempdir().unwrap();
                let storage = make_storage(dir.path());
                let ns = NamespaceId("bench-ns".into());
                let topic = TopicName("bench-topic".into());

                // Pre-write messages
                rt.block_on(async {
                    for i in 0..n as u64 {
                        let msg = make_published(i, "bench-ns", "bench-topic", 256);
                        storage.append(&ns, &topic, &msg).await.unwrap();
                    }
                    storage.flush(&ns, &topic).await.unwrap();
                });

                b.iter(|| {
                    rt.block_on(storage.replay(&ns, &topic)).unwrap();
                });
            },
        );
    }

    group.finish();
}

fn flush_latency(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();
    let mut group = c.benchmark_group("storage_flush");

    let dir = tempfile::tempdir().unwrap();
    let storage = make_storage(dir.path());
    let ns = NamespaceId("bench-ns".into());
    let topic = TopicName("bench-topic".into());

    // Write some data first
    rt.block_on(async {
        for i in 0..1_000u64 {
            let msg = make_published(i, "bench-ns", "bench-topic", 256);
            storage.append(&ns, &topic, &msg).await.unwrap();
        }
    });

    let mut write_id = 1_000u64;
    group.bench_function("flush_after_writes", |b| {
        b.iter(|| {
            // Write a few messages then flush
            rt.block_on(async {
                for _ in 0..10 {
                    let msg = make_published(write_id, "bench-ns", "bench-topic", 256);
                    write_id += 1;
                    storage.append(&ns, &topic, &msg).await.unwrap();
                }
                storage.flush(&ns, &topic).await.unwrap();
            });
        });
    });

    group.finish();
}

fn offset_at_timestamp(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();
    let mut group = c.benchmark_group("offset_at_timestamp");
    group.sample_size(10);

    // Messages per segment ≈ 5-6 with 512-byte max and ~100-byte records.
    // To span N segments we write N * 6 messages.
    for num_segments in [1, 10, 50] {
        let msg_count: u64 = num_segments * 6;

        group.bench_with_input(
            BenchmarkId::new("segments", num_segments),
            &msg_count,
            |b, &n| {
                let dir = tempfile::tempdir().unwrap();
                let data_dir = dir.path().join("data");
                let base_ts: u64 = 1_700_000_000_000;
                let ns = NamespaceId("bench-ns".into());
                let topic = TopicName("bench-topic".into());

                // Write messages with controlled monotonic timestamps, then
                // shutdown + re-init so segment indexes are built and cached.
                let storage = rt.block_on(async {
                    let storage = Arc::new(FsStorage::with_segment_max_bytes(&data_dir, 512));
                    storage.init().await.unwrap();
                    for i in 0..n {
                        let msg = make_published_with_ts(
                            i,
                            "bench-ns",
                            "bench-topic",
                            64,
                            base_ts + i * 1000,
                        );
                        storage.append(&ns, &topic, &msg).await.unwrap();
                    }
                    storage.shutdown().await.unwrap();

                    let storage = Arc::new(FsStorage::with_segment_max_bytes(&data_dir, 512));
                    storage.init().await.unwrap();
                    storage
                });

                // Target timestamp in the middle of the range
                let target_ts = base_ts + (n / 2) * 1000;

                b.iter(|| {
                    rt.block_on(storage.offset_at_timestamp(&ns, &topic, target_ts))
                        .unwrap();
                });
            },
        );
    }

    group.finish();
}

criterion_group!(
    benches,
    append_single_topic,
    append_cross_topic,
    replay_throughput,
    flush_latency,
    offset_at_timestamp,
);
criterion_main!(benches);
