use std::sync::Arc;

use criterion::{Criterion, Throughput, black_box, criterion_group, criterion_main};
use pd_proto_tso::{KeyspaceGroup, RequestHeader, TsoRequest};
use pd_service_tso::oracle::OracleConfig;
use pd_service_tso::{
    AllocatorConfig, InMemoryEtcdStore, KeyspaceGroupManager, TsoService,
    benchmark_forwarding_merge_and_split,
};

fn build_service() -> TsoService {
    let store = Arc::new(InMemoryEtcdStore::new());
    let allocator = AllocatorConfig {
        lease_ms: 1_000_000_000,
        oracle: OracleConfig {
            save_interval_ms: 1_000_000_000,
            update_physical_interval_ms: 1,
            ..Default::default()
        },
    };
    let mut manager = KeyspaceGroupManager::new(1, store, allocator, "node");
    manager
        .register_group(KeyspaceGroup {
            id: 0,
            user_kind: "basic".to_string(),
            split_state: None,
            merge_state: None,
            members: Vec::new(),
            keyspaces: vec![0],
        })
        .expect("register group");
    manager
        .campaign_group_primary(0, 10_000)
        .expect("campaign primary");
    TsoService::new(1, manager)
}

fn bench_routing_local_tso_count_1(c: &mut Criterion) {
    let service = build_service();
    let now_ms = 20_000_i64;
    c.bench_function("routing_local_tso_count_1", |b| {
        b.iter(|| {
            let response = service
                .tso(
                    now_ms,
                    TsoRequest {
                        header: RequestHeader {
                            cluster_id: 1,
                            keyspace_id: 0,
                            keyspace_group_id: 0,
                            sender_id: 0,
                        },
                        count: 1,
                        dc_location: String::new(),
                    },
                )
                .expect("tso response");
            black_box(response.timestamp.logical);
        });
    });
}

fn bench_routing_local_tso_count_64(c: &mut Criterion) {
    let service = build_service();
    let now_ms = 30_000_i64;
    c.bench_function("routing_local_tso_count_64", |b| {
        b.iter(|| {
            let response = service
                .tso(
                    now_ms,
                    TsoRequest {
                        header: RequestHeader {
                            cluster_id: 1,
                            keyspace_id: 0,
                            keyspace_group_id: 0,
                            sender_id: 0,
                        },
                        count: 64,
                        dc_location: String::new(),
                    },
                )
                .expect("tso response");
            black_box(response.timestamp.logical);
        });
    });
}

fn bench_forwarding_merge_and_split_32(c: &mut Criterion) {
    let counts = vec![1_u32; 32];
    let mut group = c.benchmark_group("forwarding");
    group.throughput(Throughput::Elements(counts.len() as u64));
    group.bench_function("forwarding_merge_and_split_32", |b| {
        b.iter(|| {
            let responses = benchmark_forwarding_merge_and_split(black_box(counts.as_slice()))
                .expect("forwarding merge+split response");
            black_box(
                responses
                    .last()
                    .and_then(|response| response.timestamp.as_ref())
                    .map(|timestamp| timestamp.logical)
                    .unwrap_or_default(),
            );
        });
    });
    group.finish();
}

fn criterion_benchmark(c: &mut Criterion) {
    bench_routing_local_tso_count_1(c);
    bench_routing_local_tso_count_64(c);
    bench_forwarding_merge_and_split_32(c);
}

criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);