use std::sync::Arc;
use std::time::Duration;

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

use fafb_q::config::{ServerConfig, shared_config};
use fafb_q::namespace::{DefaultNamespaceManager, ManagerConfig, NamespaceManager};
use fafb_q::queue::{EngineConfig, InMemoryQueueEngine};
use fafb_q::storage::mock::MockStorage;
use fafb_q::types::{
    ClientId, GroupName, MessageEnvelope, MessageId, NamespaceId, PublishedMessage, SubscribeStart,
    TopicName, now_millis,
};

fn make_envelope(id: u64, ns: &str, topic: &str, payload_size: usize) -> MessageEnvelope {
    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,
    }
}

fn make_engine() -> Arc<InMemoryQueueEngine<MockStorage>> {
    Arc::new(InMemoryQueueEngine::new(
        Arc::new(MockStorage::new()),
        EngineConfig::default(),
    ))
}

fn make_manager(
    engine: Arc<InMemoryQueueEngine<MockStorage>>,
    ns_ids: Vec<NamespaceId>,
) -> Arc<DefaultNamespaceManager<InMemoryQueueEngine<MockStorage>>> {
    Arc::new(DefaultNamespaceManager::new(
        engine,
        ns_ids,
        shared_config(&ServerConfig::default()),
        ManagerConfig::default(),
    ))
}

const RECV_TIMEOUT: Duration = Duration::from_millis(500);

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

    let ns = NamespaceId("bench-ns".into());
    let topic = TopicName("bench-topic".into());

    for num_subscribers in [1, 10, 100] {
        group.bench_with_input(
            BenchmarkId::new("subscribers", num_subscribers),
            &num_subscribers,
            |b, &n| {
                let engine = make_engine();
                let manager = make_manager(Arc::clone(&engine), vec![ns.clone()]);

                // Subscribe N clients
                let mut receivers: Vec<tokio::sync::mpsc::Receiver<PublishedMessage>> =
                    Vec::with_capacity(n);
                rt.block_on(async {
                    for i in 0..n {
                        let result = manager
                            .subscribe(
                                &ns,
                                &topic,
                                ClientId(i as u64 + 1),
                                None,
                                SubscribeStart::Head,
                            )
                            .await
                            .unwrap();
                        receivers.push(result.rx);
                    }
                });

                let mut msg_id = 0u64;
                b.iter(|| {
                    rt.block_on(async {
                        let env = make_envelope(msg_id, "bench-ns", "bench-topic", 256);
                        msg_id += 1;
                        manager.publish(&ns, &topic, env).await.unwrap();

                        // Block on channel recv with timeout — measures actual delivery latency
                        for rx in &mut receivers {
                            let _ = tokio::time::timeout(RECV_TIMEOUT, rx.recv()).await;
                        }
                    });
                });

                rt.block_on(manager.shutdown()).ok();
            },
        );
    }

    group.finish();
}

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

    let ns = NamespaceId("bench-ns".into());
    let topic = TopicName("bench-topic".into());
    let group_name = GroupName("bench-group".into());

    for num_members in [1, 10] {
        group.bench_with_input(
            BenchmarkId::new("members", num_members),
            &num_members,
            |b, &n| {
                let engine = make_engine();
                let manager = make_manager(Arc::clone(&engine), vec![ns.clone()]);

                // Subscribe N group members
                let mut receivers: Vec<tokio::sync::mpsc::Receiver<PublishedMessage>> =
                    Vec::with_capacity(n);
                rt.block_on(async {
                    for i in 0..n {
                        let result = manager
                            .subscribe(
                                &ns,
                                &topic,
                                ClientId(i as u64 + 1),
                                Some(group_name.clone()),
                                SubscribeStart::Head,
                            )
                            .await
                            .unwrap();
                        receivers.push(result.rx);
                    }
                });

                let mut msg_id = 0u64;
                b.iter(|| {
                    rt.block_on(async {
                        // Publish N messages (so each member gets ~1)
                        for _ in 0..n {
                            let env = make_envelope(msg_id, "bench-ns", "bench-topic", 256);
                            msg_id += 1;
                            manager.publish(&ns, &topic, env).await.unwrap();
                        }

                        // Round-robin delivers to arbitrary members — poll all with
                        // short timeouts to avoid blocking on the wrong receiver.
                        let mut received = 0;
                        let short = Duration::from_millis(1);
                        while received < n {
                            for rx in &mut receivers {
                                if let Ok(Some(_)) = tokio::time::timeout(short, rx.recv()).await {
                                    received += 1;
                                }
                                if received >= n {
                                    break;
                                }
                            }
                        }
                    });
                });

                rt.block_on(manager.shutdown()).ok();
            },
        );
    }

    group.finish();
}

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

    let ns = NamespaceId("bench-ns".into());
    let topic = TopicName("bench-topic".into());

    // Publish without any subscribers — measures pure namespace manager overhead
    group.bench_function("no_subscribers", |b| {
        let engine = make_engine();
        let manager = make_manager(Arc::clone(&engine), vec![ns.clone()]);
        let mut msg_id = 0u64;

        b.iter(|| {
            let env = make_envelope(msg_id, "bench-ns", "bench-topic", 256);
            msg_id += 1;
            rt.block_on(manager.publish(&ns, &topic, env)).unwrap();
        });

        rt.block_on(manager.shutdown()).ok();
    });

    group.finish();
}

criterion_group!(
    benches,
    individual_fanout,
    consumer_group_delivery,
    publish_throughput,
);
criterion_main!(benches);
