use std::sync::Arc;

use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
use iroh::{SecretKey, address_lookup::pkarr::PkarrRelayClient, endpoint_info::EndpointInfo};
use iroh_dns_server::{ZoneStore, config::Config, metrics::Metrics, server::Server};
use n0_error::Result;
use rand_chacha::rand_core::SeedableRng;
use tokio::runtime::Runtime;

const LOCALHOST_PKARR: &str = "http://localhost:8080/pkarr";

async fn start_dns_server(config: Config) -> Result<Server> {
    let metrics = Arc::new(Metrics::default());
    let store = ZoneStore::persistent(
        config.signed_packet_store_path()?,
        Default::default(),
        metrics.clone(),
    )?;
    Server::spawn(config, store, metrics).await
}

fn benchmark_dns_server(c: &mut Criterion) {
    let mut group = c.benchmark_group("dns_server_writes");
    group.sample_size(10);
    for iters in [10_u64, 100_u64, 250_u64, 1000_u64].iter() {
        group.throughput(Throughput::Elements(*iters));
        group.bench_with_input(BenchmarkId::from_parameter(iters), iters, |b, &iters| {
            b.iter(|| {
                let rt = Runtime::new().unwrap();
                rt.block_on(async move {
                    let config = Config::load("./config.dev.toml").await.unwrap();
                    let server = start_dns_server(config).await.unwrap();

                    let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(42);
                    let secret_key = SecretKey::generate(&mut rng);
                    let endpoint_id = secret_key.public();

                    let pkarr_relay = LOCALHOST_PKARR.parse().expect("valid url");
                    let pkarr = PkarrRelayClient::new(pkarr_relay);
                    let relay_url = "http://localhost:8080".parse().unwrap();
                    let endpoint_info =
                        EndpointInfo::new(endpoint_id).with_relay_url(Some(relay_url));
                    let signed_packet = endpoint_info
                        .to_pkarr_signed_packet(&secret_key, 30)
                        .unwrap();

                    let start = std::time::Instant::now();
                    for _ in 0..iters {
                        pkarr.publish(&signed_packet).await.unwrap();
                    }
                    let duration = start.elapsed();

                    server.shutdown().await.unwrap();

                    duration
                })
            });
        });
    }
}

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

Neighbours