// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use benchmarks::*;
use tls_harness::cohort::S2NConnection;

use criterion::{
    criterion_group, criterion_main, measurement::WallTime, BatchSize, BenchmarkGroup, Criterion,
};
use tls_harness::{harness::TlsInfo, SigType, TlsConnection};

fn bench_handshake_pair<T>(bench_group: &mut BenchmarkGroup<WallTime>, sig_type: SigType)
where
    T: TlsConnection + TlsInfo,
    T::Config: TlsBenchConfig,
{
    // generate all harnesses (TlsConnPair structs) beforehand so that benchmarks
    // only include negotiation and not config/connection initialization
    for handshake in [HandshakeType::Resumption, HandshakeType::ServerAuth] {
        bench_group.bench_function(format!("{:?}-{}", handshake, T::name()), |b| {
            b.iter_batched_ref(
                || {
                    new_bench_pair::<T, T>(
                        CryptoConfig::new(CipherSuite::default(), KXGroup::default(), sig_type),
                        handshake,
                    )
                },
                |conn_pair_res| {
                    if let Ok(conn_pair) = conn_pair_res {
                        let _ = conn_pair.handshake();
                    }
                },
                BatchSize::SmallInput,
            )
        });
    }
}

fn bench_handshake_server_1rtt<T>(bench_group: &mut BenchmarkGroup<WallTime>, sig_type: SigType)
where
    T: TlsConnection + TlsInfo,
    T::Config: TlsBenchConfig,
{
    for handshake in [HandshakeType::Resumption, HandshakeType::ServerAuth] {
        bench_group.bench_function(format!("{:?}-{}", handshake, T::name()), |b| {
            b.iter_batched_ref(
                || {
                    let mut pair = new_bench_pair::<T, T>(
                        CryptoConfig::new(CipherSuite::default(), KXGroup::default(), sig_type),
                        handshake,
                    )
                    .unwrap();
                    pair.client_mut().handshake().unwrap();
                    pair
                },
                |pair| {
                    // this represents the work that the server does during the
                    // first RTT
                    pair.server_mut().handshake().unwrap();
                },
                BatchSize::SmallInput,
            )
        });
    }
}

/// This benchmark compares resumption savings across a single implementation.
/// E.g. "how much faster is session resumption than a full handshake for
/// s2n-tls?".
pub fn bench_resumption(c: &mut Criterion) {
    // compare resumption savings across both client and server
    for sig_type in [SigType::Rsa2048, SigType::Ecdsa256] {
        let mut bench_group = c.benchmark_group(format!("resumption-pair-{sig_type:?}"));
        bench_handshake_pair::<S2NConnection>(&mut bench_group, sig_type);
    }

    // only look at resumption savings for the server, specifically the work
    // done in the first rtt.
    for sig_type in [SigType::Rsa2048, SigType::Ecdsa384] {
        let mut bench_group = c.benchmark_group(format!("resumption-server-1rtt-{sig_type:?}"));
        bench_handshake_server_1rtt::<S2NConnection>(&mut bench_group, sig_type);
    }
}

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