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

use rustls::{cipher_suite as ciphers, quic, CipherSuite, SupportedCipherSuite};
use s2n_quic_core::crypto::{self, tls, CryptoError, HeaderProtectionMask, Key};

pub struct PacketKey {
    key: quic::PacketKey,
    cipher_suite: tls::CipherSuite,
}

impl PacketKey {
    pub(crate) fn new(
        keys: quic::DirectionalKeys,
        cipher_suite: tls::CipherSuite,
    ) -> (Self, HeaderProtectionKey) {
        (
            Self {
                key: keys.packet,
                cipher_suite,
            },
            HeaderProtectionKey(keys.header),
        )
    }
}

impl crypto::Key for PacketKey {
    fn decrypt(
        &self,
        packet_number: u64,
        header: &[u8],
        payload: &mut [u8],
    ) -> Result<(), CryptoError> {
        match self.key.decrypt_in_place(packet_number, header, payload) {
            Ok(_tag) => Ok(()),
            Err(_) => Err(CryptoError::DECRYPT_ERROR),
        }
    }

    fn encrypt(
        &self,
        packet_number: u64,
        header: &[u8],
        payload: &mut [u8],
    ) -> Result<(), CryptoError> {
        let (payload, tag_storage) = payload.split_at_mut(payload.len() - self.tag_len());
        match self.key.encrypt_in_place(packet_number, header, payload) {
            Ok(tag) => {
                tag_storage.copy_from_slice(tag.as_ref());
                Ok(())
            }
            Err(_) => Err(CryptoError::INTERNAL_ERROR),
        }
    }

    fn tag_len(&self) -> usize {
        self.key.tag_len()
    }

    fn aead_confidentiality_limit(&self) -> u64 {
        self.key.confidentiality_limit()
    }

    fn aead_integrity_limit(&self) -> u64 {
        self.key.integrity_limit()
    }

    fn cipher_suite(&self) -> tls::CipherSuite {
        self.cipher_suite
    }
}

impl crypto::ZeroRttKey for PacketKey {}

pub struct PacketKeys {
    sealer: PacketKey,
    opener: PacketKey,
}

impl PacketKeys {
    pub(crate) fn new(keys: quic::Keys, cipher_suite: CipherSuite) -> (Self, HeaderProtectionKeys) {
        let quic::Keys { local, remote } = keys;

        let cipher_suite = match cipher_suite {
            CipherSuite::TLS13_AES_128_GCM_SHA256 => tls::CipherSuite::TLS_AES_128_GCM_SHA256,
            CipherSuite::TLS13_AES_256_GCM_SHA384 => tls::CipherSuite::TLS_AES_256_GCM_SHA384,
            CipherSuite::TLS13_CHACHA20_POLY1305_SHA256 => {
                tls::CipherSuite::TLS_CHACHA20_POLY1305_SHA256
            }
            _ => tls::CipherSuite::Unknown,
        };

        let (sealer_packet, sealer_header) = PacketKey::new(local, cipher_suite);
        let (opener_packet, opener_header) = PacketKey::new(remote, cipher_suite);

        let key = Self {
            sealer: sealer_packet,
            opener: opener_packet,
        };

        let header_key = HeaderProtectionKeys {
            sealer: sealer_header,
            opener: opener_header,
        };

        (key, header_key)
    }
}

impl crypto::Key for PacketKeys {
    fn decrypt(
        &self,
        packet_number: u64,
        header: &[u8],
        payload: &mut [u8],
    ) -> Result<(), CryptoError> {
        self.opener.decrypt(packet_number, header, payload)
    }

    fn encrypt(
        &self,
        packet_number: u64,
        header: &[u8],
        payload: &mut [u8],
    ) -> Result<(), CryptoError> {
        self.sealer.encrypt(packet_number, header, payload)
    }

    fn tag_len(&self) -> usize {
        self.sealer.tag_len()
    }

    fn aead_confidentiality_limit(&self) -> u64 {
        self.sealer.aead_confidentiality_limit()
    }

    fn aead_integrity_limit(&self) -> u64 {
        self.sealer.aead_integrity_limit()
    }

    fn cipher_suite(&self) -> tls::CipherSuite {
        self.sealer.cipher_suite()
    }
}

impl crypto::HandshakeKey for PacketKeys {}

pub struct HeaderProtectionKey(quic::HeaderProtectionKey);

impl HeaderProtectionKey {
    /// Returns the header protection mask for the given ciphertext sample
    ///
    /// Rustls API applies the header protection rather than returning
    /// the mask. This method exists for extracting the mask from rustls by calling
    /// the `encrypt_in_place` api and reversing the operation it applies.
    ///
    /// The primary motivation for extracting the mask from rustls is to maintain API
    /// compatibility with other tls providers (s2n-tls) and have one common code
    /// for applying and removing header protection.
    fn get_mask(&self, ciphertext_sample: &[u8]) -> HeaderProtectionMask {
        let mut mask = HeaderProtectionMask::default();

        // tell rustls we need all of the packet number mask by setting the pn length to 4
        let tag = 0b0000_0011;
        mask[0] = tag;

        let (packet_tag, packet_number) = mask.split_first_mut().unwrap();
        self.0
            .encrypt_in_place(ciphertext_sample, packet_tag, packet_number)
            .unwrap();

        // rustls XORs the mask with the original tag so undo that operation
        mask[0] ^= tag;

        mask
    }
}

impl crypto::HeaderKey for HeaderProtectionKey {
    fn opening_header_protection_mask(&self, ciphertext_sample: &[u8]) -> HeaderProtectionMask {
        self.get_mask(ciphertext_sample)
    }

    fn opening_sample_len(&self) -> usize {
        self.0.sample_len()
    }

    fn sealing_header_protection_mask(&self, ciphertext_sample: &[u8]) -> HeaderProtectionMask {
        self.get_mask(ciphertext_sample)
    }

    fn sealing_sample_len(&self) -> usize {
        self.0.sample_len()
    }
}

impl crypto::ZeroRttHeaderKey for HeaderProtectionKey {}

pub struct HeaderProtectionKeys {
    sealer: HeaderProtectionKey,
    opener: HeaderProtectionKey,
}

impl crypto::HeaderKey for HeaderProtectionKeys {
    fn opening_header_protection_mask(&self, ciphertext_sample: &[u8]) -> HeaderProtectionMask {
        self.opener
            .opening_header_protection_mask(ciphertext_sample)
    }

    fn opening_sample_len(&self) -> usize {
        self.opener.opening_sample_len()
    }

    fn sealing_header_protection_mask(&self, ciphertext_sample: &[u8]) -> HeaderProtectionMask {
        self.sealer
            .sealing_header_protection_mask(ciphertext_sample)
    }

    fn sealing_sample_len(&self) -> usize {
        self.sealer.sealing_sample_len()
    }
}

impl crypto::HandshakeHeaderKey for HeaderProtectionKeys {}
impl crypto::OneRttHeaderKey for HeaderProtectionKeys {}

pub struct OneRttKey {
    key: PacketKeys,
    secrets: quic::Secrets,
}

impl OneRttKey {
    pub(crate) fn new(
        keys: quic::Keys,
        secrets: quic::Secrets,
        cipher_suite: CipherSuite,
    ) -> (Self, HeaderProtectionKeys) {
        let (key, header_key) = PacketKeys::new(keys, cipher_suite);
        let key = Self { key, secrets };
        (key, header_key)
    }
}

impl crypto::Key for OneRttKey {
    fn decrypt(
        &self,
        packet_number: u64,
        header: &[u8],
        payload: &mut [u8],
    ) -> Result<(), CryptoError> {
        self.key.decrypt(packet_number, header, payload)
    }

    fn encrypt(
        &self,
        packet_number: u64,
        header: &[u8],
        payload: &mut [u8],
    ) -> Result<(), CryptoError> {
        self.key.encrypt(packet_number, header, payload)
    }

    fn tag_len(&self) -> usize {
        self.key.tag_len()
    }

    fn aead_confidentiality_limit(&self) -> u64 {
        self.key.aead_confidentiality_limit()
    }

    fn aead_integrity_limit(&self) -> u64 {
        self.key.aead_integrity_limit()
    }

    fn cipher_suite(&self) -> tls::CipherSuite {
        self.key.cipher_suite()
    }
}

impl crypto::OneRttKey for OneRttKey {
    fn derive_next_key(&self) -> Self {
        let cipher_suite = self.cipher_suite();
        let mut secrets = self.secrets.clone();
        let quic::PacketKeySet { local, remote } = secrets.next_packet_keys();
        Self {
            key: PacketKeys {
                sealer: PacketKey {
                    key: local,
                    cipher_suite,
                },
                opener: PacketKey {
                    key: remote,
                    cipher_suite,
                },
            },
            secrets,
        }
    }

    fn update_sealer_pmtu(&mut self, _pmtu: u16) {
        // rustls doesn't have any ptmu specialization
    }

    fn update_opener_pmtu(&mut self, _pmtu: u16) {
        // rustls doesn't have any ptmu specialization
    }
}

//= https://www.rfc-editor.org/rfc/rfc9001#section-5.3
//# A cipher suite MUST NOT be
//# negotiated unless a header protection scheme is defined for the
//# cipher suite.
// All of the cipher_suites from the current exported list have HP schemes for QUIC
pub static DEFAULT_CIPHERSUITES: &[SupportedCipherSuite] = &[
    ciphers::TLS13_AES_128_GCM_SHA256,
    ciphers::TLS13_AES_256_GCM_SHA384,
    ciphers::TLS13_CHACHA20_POLY1305_SHA256,
];

#[test]
fn test_default_cipher_suites() {
    insta::assert_debug_snapshot!("default_cipher_suites", DEFAULT_CIPHERSUITES);
}


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

use crate::{certificate, encode_transport_parameters, session::Session};
use core::convert::TryFrom;
use rustls::{quic, ClientConfig};
use s2n_codec::EncoderValue;
use s2n_quic_core::{application::ServerName, crypto::tls};
use std::sync::Arc;

#[derive(Clone)]
pub struct Client {
    config: Arc<ClientConfig>,
}

impl Client {
    pub fn new(config: ClientConfig) -> Self {
        Self {
            config: Arc::new(config),
        }
    }

    pub fn builder() -> Builder {
        Builder::new()
    }
}

impl Default for Client {
    fn default() -> Self {
        // TODO this will currently panic since there is no default root cert
        Self::builder()
            .build()
            .expect("could not create default client")
    }
}

impl From<ClientConfig> for Client {
    fn from(config: ClientConfig) -> Self {
        Self::new(config)
    }
}

impl From<Arc<ClientConfig>> for Client {
    fn from(config: Arc<ClientConfig>) -> Self {
        Self { config }
    }
}

impl tls::Endpoint for Client {
    type Session = Session;

    fn new_server_session<Params: EncoderValue>(
        &mut self,
        _transport_parameters: &Params,
    ) -> Self::Session {
        panic!("cannot create a server session from a client config");
    }

    fn new_client_session<Params: EncoderValue>(
        &mut self,
        transport_parameters: &Params,
        server_name: ServerName,
    ) -> Self::Session {
        use quic::ClientQuicExt;

        //= https://www.rfc-editor.org/rfc/rfc9001#section-8.2
        //# Endpoints MUST send the quic_transport_parameters extension;
        let transport_parameters = encode_transport_parameters(transport_parameters);

        let rustls_server_name =
            rustls::ServerName::try_from(server_name.as_ref()).expect("invalid server name");

        let session = rustls::ClientConnection::new_quic(
            self.config.clone(),
            crate::QUIC_VERSION,
            rustls_server_name,
            transport_parameters,
        )
        .expect("could not create rustls client session");

        Session::new(session.into(), Some(server_name))
    }

    fn max_tag_length(&self) -> usize {
        s2n_quic_crypto::MAX_TAG_LEN
    }
}

pub struct Builder {
    cert_store: rustls::RootCertStore,
    application_protocols: Vec<Vec<u8>>,
    key_log: Option<Arc<dyn rustls::KeyLog>>,
}

impl Default for Builder {
    fn default() -> Self {
        Self::new()
    }
}

impl Builder {
    pub fn new() -> Self {
        Self {
            cert_store: rustls::RootCertStore::empty(),
            application_protocols: vec![b"h3".to_vec()],
            key_log: None,
        }
    }

    pub fn with_certificate<C: certificate::IntoCertificate>(
        mut self,
        certificate: C,
    ) -> Result<Self, rustls::Error> {
        let certificates = certificate.into_certificate()?;
        let root_certificate = certificates.0.get(0).ok_or_else(|| {
            rustls::Error::General("Certificate chain needs to have at least one entry".to_string())
        })?;
        self.cert_store
            .add(root_certificate)
            .map_err(|err| rustls::Error::General(err.to_string()))?;
        Ok(self)
    }

    pub fn with_max_cert_chain_depth(self, len: u16) -> Result<Self, rustls::Error> {
        // TODO is there a way to configure this?
        let _ = len;
        Ok(self)
    }

    pub fn with_application_protocols<P: Iterator<Item = I>, I: AsRef<[u8]>>(
        mut self,
        protocols: P,
    ) -> Result<Self, rustls::Error> {
        self.application_protocols = protocols.map(|p| p.as_ref().to_vec()).collect();
        Ok(self)
    }

    pub fn with_key_logging(mut self) -> Result<Self, rustls::Error> {
        self.key_log = Some(Arc::new(rustls::KeyLogFile::new()));
        Ok(self)
    }

    pub fn build(self) -> Result<Client, rustls::Error> {
        // TODO load system root store?
        if self.cert_store.is_empty() {
            //= https://www.rfc-editor.org/rfc/rfc9001#section-4.4
            //# A client MUST authenticate the identity of the server.
            return Err(rustls::Error::General(
                "missing trusted root certificate(s)".to_string(),
            ));
        }

        let mut config = ClientConfig::builder()
            .with_cipher_suites(crate::cipher_suite::DEFAULT_CIPHERSUITES)
            .with_safe_default_kx_groups()
            .with_protocol_versions(crate::PROTOCOL_VERSIONS)?
            .with_root_certificates(self.cert_store)
            .with_no_client_auth();

        config.max_fragment_size = None;
        config.alpn_protocols = self.application_protocols;

        if let Some(key_log) = self.key_log {
            config.key_log = key_log;
        }

        Ok(Client::new(config))
    }
}


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

use rustls::Error;

pub fn reason(error: rustls::Error) -> &'static str {
    match error {
        Error::InappropriateMessage { .. } => "received unexpected message",
        Error::InappropriateHandshakeMessage { .. } => "received unexpected handshake message",
        Error::CorruptMessage | Error::CorruptMessagePayload(_) => "received corrupt message",
        Error::NoCertificatesPresented => "peer sent no certificates",
        Error::UnsupportedNameType => "unsupported name type",
        Error::DecryptError => "cannot decrypt peer's message",
        Error::EncryptError => "cannot encrypt local message",
        Error::PeerIncompatibleError(_) => "peer is incompatible",
        Error::PeerMisbehavedError(_) => "peer misbehaved",
        Error::AlertReceived(_) => "received fatal alert",
        Error::InvalidCertificateEncoding => "invalid certificate encoding",
        Error::InvalidCertificateSignatureType => "invalid certificate signature type",
        Error::InvalidCertificateSignature => "invalid certificate signature",
        Error::InvalidCertificateData(_) => "invalid certificate data",
        Error::InvalidSct(_) => "invalid certificate timestamp",
        Error::FailedToGetCurrentTime => "failed to get current time",
        Error::FailedToGetRandomBytes => "failed to get random bytes",
        Error::HandshakeNotComplete => "handshake not complete",
        Error::PeerSentOversizedRecord => "peer sent excess record size",
        Error::NoApplicationProtocol => "peer doesn't support any known protocol",
        Error::BadMaxFragmentSize => "bad max fragment size",
        Error::General(_) => "unexpected error",
        // rustls may add a new variant in the future that breaks us so do a wildcard
        #[allow(unreachable_patterns)]
        _ => "unexpected error",
    }
}


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

#![forbid(unsafe_code)]

pub use rustls::{self, Certificate, PrivateKey};

mod cipher_suite;
mod error;
mod session;

pub mod certificate;
pub mod client;
pub mod server;

pub use client::Client;
pub use server::Server;

//= https://www.rfc-editor.org/rfc/rfc9001#section-4.2
//# Clients MUST NOT offer TLS versions older than 1.3.
static PROTOCOL_VERSIONS: &[&rustls::SupportedProtocolVersion] = &[&rustls::version::TLS13];

/// The supported version of quic
const QUIC_VERSION: rustls::quic::Version = rustls::quic::Version::V1;

/// Encodes transport parameters into a byte vec
pub(crate) fn encode_transport_parameters<Params: s2n_codec::EncoderValue>(
    params: &Params,
) -> Vec<u8> {
    let len = params.encoding_size();
    let mut buffer = vec![0; len];
    params.encode(&mut s2n_codec::EncoderBuffer::new(&mut buffer));
    buffer
}

#[test]
fn client_server_test() {
    use s2n_quic_core::crypto::tls::{self, testing::certificates::*};

    let mut client = client::Builder::new()
        .with_certificate(CERT_PEM)
        .unwrap()
        .build()
        .unwrap();

    let mut server = server::Builder::new()
        .with_certificate(CERT_PEM, KEY_PEM)
        .unwrap()
        .build()
        .unwrap();

    let mut pair = tls::testing::Pair::new(&mut server, &mut client, "localhost".into());

    while pair.is_handshaking() {
        pair.poll(None).unwrap();
    }

    pair.finish();
}


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

use crate::{certificate, encode_transport_parameters, session::Session};
use rustls::{quic, ServerConfig};
use s2n_codec::EncoderValue;
use s2n_quic_core::{application::ServerName, crypto::tls};
use std::sync::Arc;

#[derive(Clone)]
pub struct Server {
    config: Arc<ServerConfig>,
}

impl Server {
    pub fn new(config: ServerConfig) -> Self {
        Self {
            config: Arc::new(config),
        }
    }

    pub fn builder() -> Builder {
        Builder::new()
    }
}

impl Default for Server {
    fn default() -> Self {
        Self::builder()
            .build()
            .expect("could not create default server")
    }
}

impl From<ServerConfig> for Server {
    fn from(config: ServerConfig) -> Self {
        Self::new(config)
    }
}

impl From<Arc<ServerConfig>> for Server {
    fn from(config: Arc<ServerConfig>) -> Self {
        Self { config }
    }
}

impl tls::Endpoint for Server {
    type Session = Session;

    fn new_server_session<Params: EncoderValue>(
        &mut self,
        transport_parameters: &Params,
    ) -> Self::Session {
        use quic::ServerQuicExt;

        //= https://www.rfc-editor.org/rfc/rfc9001#section-8.2
        //# Endpoints MUST send the quic_transport_parameters extension;
        let transport_parameters = encode_transport_parameters(transport_parameters);

        let session = rustls::ServerConnection::new_quic(
            self.config.clone(),
            crate::QUIC_VERSION,
            transport_parameters,
        )
        .expect("could not create rustls server session");

        Session::new(session.into(), None)
    }

    fn new_client_session<Params: EncoderValue>(
        &mut self,
        _transport_parameters: &Params,
        _sni: ServerName,
    ) -> Self::Session {
        panic!("cannot create a client session from a server config");
    }

    fn max_tag_length(&self) -> usize {
        s2n_quic_crypto::MAX_TAG_LEN
    }
}

pub struct Builder {
    cert_resolver: Option<Arc<dyn rustls::server::ResolvesServerCert>>,
    application_protocols: Vec<Vec<u8>>,
    key_log: Option<Arc<dyn rustls::KeyLog>>,
}

impl Default for Builder {
    fn default() -> Self {
        Self::new()
    }
}

impl Builder {
    pub fn new() -> Self {
        Self {
            cert_resolver: None,
            application_protocols: vec![b"h3".to_vec()],
            key_log: None,
        }
    }

    pub fn with_certificate<C: certificate::IntoCertificate, PK: certificate::IntoPrivateKey>(
        self,
        certificate: C,
        private_key: PK,
    ) -> Result<Self, rustls::Error> {
        let certificate = certificate.into_certificate()?;
        let private_key = private_key.into_private_key()?;
        let resolver = AlwaysResolvesChain::new(certificate, private_key)?;
        let resolver = Arc::new(resolver);
        self.with_cert_resolver(resolver)
    }

    pub fn with_cert_resolver(
        mut self,
        cert_resolver: Arc<dyn rustls::server::ResolvesServerCert>,
    ) -> Result<Self, rustls::Error> {
        self.cert_resolver = Some(cert_resolver);
        Ok(self)
    }

    pub fn with_application_protocols<P: Iterator<Item = I>, I: AsRef<[u8]>>(
        mut self,
        protocols: P,
    ) -> Result<Self, rustls::Error> {
        self.application_protocols = protocols.map(|p| p.as_ref().to_vec()).collect();
        Ok(self)
    }

    pub fn with_key_logging(mut self) -> Result<Self, rustls::Error> {
        self.key_log = Some(Arc::new(rustls::KeyLogFile::new()));
        Ok(self)
    }

    pub fn build(self) -> Result<Server, rustls::Error> {
        let builder = ServerConfig::builder()
            .with_cipher_suites(crate::cipher_suite::DEFAULT_CIPHERSUITES)
            .with_safe_default_kx_groups()
            .with_protocol_versions(crate::PROTOCOL_VERSIONS)?
            .with_no_client_auth();

        let mut config = if let Some(cert_resolver) = self.cert_resolver {
            builder.with_cert_resolver(cert_resolver)
        } else {
            return Err(rustls::Error::General(
                "Missing certificate or certificate resolver".to_string(),
            ));
        };

        config.ignore_client_order = true;
        config.max_fragment_size = None;
        config.alpn_protocols = self.application_protocols;

        if let Some(key_log) = self.key_log {
            config.key_log = key_log;
        }

        Ok(Server::new(config))
    }
}

struct AlwaysResolvesChain(Arc<rustls::sign::CertifiedKey>);

impl AlwaysResolvesChain {
    fn new(
        chain: certificate::Certificate,
        priv_key: certificate::PrivateKey,
    ) -> Result<Self, rustls::Error> {
        let key = rustls::sign::any_supported_type(&priv_key.0)
            .map_err(|_| rustls::Error::General("invalid private key".into()))?;
        Ok(Self(Arc::new(rustls::sign::CertifiedKey::new(
            chain.0, key,
        ))))
    }
}

impl rustls::server::ResolvesServerCert for AlwaysResolvesChain {
    fn resolve(
        &self,
        _client_hello: rustls::server::ClientHello,
    ) -> Option<Arc<rustls::sign::CertifiedKey>> {
        Some(Arc::clone(&self.0))
    }
}


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

use crate::cipher_suite::{
    HeaderProtectionKey, HeaderProtectionKeys, OneRttKey, PacketKey, PacketKeys,
};
use bytes::Bytes;
use core::{fmt, fmt::Debug, task::Poll};
use rustls::{
    quic::{self, QuicExt},
    Connection,
};
use s2n_quic_core::{
    application::ServerName,
    crypto::{self, tls, CryptoError},
    transport,
};

pub struct Session {
    connection: Connection,
    rx_phase: HandshakePhase,
    tx_phase: HandshakePhase,
    emitted_zero_rtt_keys: bool,
    emitted_handshake_complete: bool,
    emitted_server_name: bool,
    emitted_application_protocol: bool,
    server_name: Option<ServerName>,
}

impl fmt::Debug for Session {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Session")
            .field("rx_phase", &self.rx_phase)
            .field("tx_phase", &self.tx_phase)
            .finish()
    }
}

impl Session {
    pub fn new(connection: Connection, server_name: Option<ServerName>) -> Self {
        Self {
            connection,
            rx_phase: Default::default(),
            tx_phase: Default::default(),
            emitted_zero_rtt_keys: false,
            emitted_handshake_complete: false,
            emitted_server_name: false,
            emitted_application_protocol: false,
            server_name,
        }
    }

    fn receive(&mut self, crypto_data: &[u8]) -> Result<(), transport::Error> {
        self.connection
            .read_hs(crypto_data)
            .map_err(crate::error::reason)
            .map_err(|reason| {
                //= https://www.rfc-editor.org/rfc/rfc9001#section-4.8
                //# QUIC is only able to convey an alert level of "fatal".  In TLS 1.3,
                //# the only existing uses for the "warning" level are to signal
                //# connection close; see Section 6.1 of [TLS13].  As QUIC provides
                //# alternative mechanisms for connection termination and the TLS
                //# connection is only closed if an error is encountered, a QUIC endpoint
                //# MUST treat any alert from TLS as if it were at the "fatal" level.

                // According to the rustls docs, `alert` only returns fatal alerts:
                // > https://docs.rs/rustls/0.19.0/rustls/quic/trait.QuicExt#tymethod.get_alert
                // > Emit the TLS description code of a fatal alert, if one has arisen.

                self.connection
                    .alert()
                    .map(|alert| CryptoError {
                        code: alert.get_u8(),
                        reason,
                    })
                    .unwrap_or(CryptoError::INTERNAL_ERROR)
            })?;
        Ok(())
    }

    fn application_parameters(&self) -> Result<tls::ApplicationParameters, transport::Error> {
        //= https://www.rfc-editor.org/rfc/rfc9001#section-8.2
        //# endpoints that
        //# receive ClientHello or EncryptedExtensions messages without the
        //# quic_transport_parameters extension MUST close the connection with an
        //# error of type 0x16d (equivalent to a fatal TLS missing_extension
        //# alert, see Section 4.8).
        let transport_parameters =
            self.connection.quic_transport_parameters().ok_or_else(|| {
                CryptoError::MISSING_EXTENSION.with_reason("Missing QUIC transport parameters")
            })?;

        Ok(tls::ApplicationParameters {
            transport_parameters,
        })
    }

    //= https://www.rfc-editor.org/rfc/rfc9001#section-8.1
    //# Unless
    //# another mechanism is used for agreeing on an application protocol,
    //# endpoints MUST use ALPN for this purpose.
    //
    //= https://www.rfc-editor.org/rfc/rfc7301#section-3.1
    //# Client                                              Server
    //#
    //#    ClientHello                     -------->       ServerHello
    //#      (ALPN extension &                               (ALPN extension &
    //#       list of protocols)                              selected protocol)
    //#                                                    [ChangeCipherSpec]
    //#                                    <--------       Finished
    //#    [ChangeCipherSpec]
    //#    Finished                        -------->
    //#    Application Data                <------->       Application Data
    fn application_protocol(&self) -> Option<&[u8]> {
        self.connection.alpn_protocol()
    }

    fn server_name(&self) -> Option<ServerName> {
        match &self.connection {
            Connection::Client(_) => self.server_name.clone(),
            Connection::Server(server) => {
                server.sni_hostname().map(|server_name| server_name.into())
            }
        }
    }

    fn zero_rtt_keys(&mut self) -> Option<quic::DirectionalKeys> {
        if self.emitted_zero_rtt_keys {
            return None;
        }

        let keys = self.connection.zero_rtt_keys()?;
        self.emitted_zero_rtt_keys = true;
        Some(keys)
    }

    /// Check and process TLS handshake complete.
    ///
    /// Upon TLS handshake complete, emit an event to notify the transport layer.
    fn poll_complete_handshake<C: tls::Context<Self>>(
        &mut self,
        context: &mut C,
    ) -> Poll<Result<(), transport::Error>> {
        if self.tx_phase == HandshakePhase::Application && !self.connection.is_handshaking() {
            // attempt to emit server_name and application_protocol events prior to completing the
            // handshake
            self.emit_events(context)?;

            // the handshake is complete!
            if !self.emitted_handshake_complete {
                self.rx_phase.transition();
                context.on_handshake_complete()?;
            }

            self.emitted_handshake_complete = true;
        }

        if self.emitted_handshake_complete {
            Poll::Ready(Ok(()))
        } else {
            Poll::Pending
        }
    }

    fn poll_impl<C: tls::Context<Self>>(
        &mut self,
        context: &mut C,
    ) -> Poll<Result<(), transport::Error>> {
        // Tracks if we have attempted to receive data at least once
        let mut has_tried_receive = false;

        loop {
            let crypto_data = match self.rx_phase {
                HandshakePhase::Initial => context.receive_initial(None),
                HandshakePhase::Handshake => context.receive_handshake(None),
                HandshakePhase::Application => context.receive_application(None),
            };

            // receive anything in the incoming buffer
            if let Some(crypto_data) = crypto_data {
                self.receive(&crypto_data)?;
            } else if has_tried_receive {
                return self.poll_complete_handshake(context);
                // If there's nothing to receive then we're done for now
            }

            if let Poll::Ready(()) = self.poll_complete_handshake(context)? {
                return Poll::Ready(Ok(()));
            }

            // mark that we tried to receive some data so we know next time we loop
            // to bail if nothing changed
            has_tried_receive = true;

            // try to pull out the early secrets, if any
            if let Some(keys) = self.zero_rtt_keys() {
                let (key, header_key) = PacketKey::new(
                    keys,
                    s2n_quic_core::crypto::tls::CipherSuite::TLS_AES_128_GCM_SHA256,
                );
                context.on_zero_rtt_keys(key, header_key, self.application_parameters()?)?;
            }

            loop {
                // make sure we can send data before pulling it out of rustls
                let can_send = match self.tx_phase {
                    HandshakePhase::Initial => context.can_send_initial(),
                    HandshakePhase::Handshake => context.can_send_handshake(),
                    HandshakePhase::Application => context.can_send_application(),
                };

                if !can_send {
                    break;
                }

                let mut transmission_buffer = vec![];

                let key_change = self.connection.write_hs(&mut transmission_buffer);

                // if we didn't upgrade the key or transmit anything then we're waiting for
                // more reads
                if key_change.is_none() && transmission_buffer.is_empty() {
                    break;
                }

                // fill the correct buffer according to the handshake phase
                match self.tx_phase {
                    HandshakePhase::Initial => context.send_initial(transmission_buffer.into()),
                    HandshakePhase::Handshake => context.send_handshake(transmission_buffer.into()),
                    HandshakePhase::Application => {
                        context.send_application(transmission_buffer.into())
                    }
                }

                if let Some(key_change) = key_change {
                    let cipher_suite = self
                        .connection
                        .negotiated_cipher_suite()
                        .expect("cipher_suite should be negotiated")
                        .suite();
                    match key_change {
                        quic::KeyChange::Handshake { keys } => {
                            let (key, header_key) = PacketKeys::new(keys, cipher_suite);

                            context.on_handshake_keys(key, header_key)?;

                            // Transition both phases to Handshake
                            self.tx_phase.transition();
                            self.rx_phase.transition();
                        }
                        quic::KeyChange::OneRtt { keys, next } => {
                            let (key, header_key) = OneRttKey::new(keys, next, cipher_suite);

                            let application_parameters = self.application_parameters()?;

                            context.on_one_rtt_keys(key, header_key, application_parameters)?;

                            // Transition the tx_phase to Application
                            // Note: the rx_phase is transitioned when the handshake is complete
                            self.tx_phase.transition();
                        }
                    }
                }
            }
        }
    }

    fn emit_events<C: tls::Context<Self>>(
        &mut self,
        context: &mut C,
    ) -> Result<(), transport::Error> {
        if !self.emitted_server_name {
            if let Some(server_name) = self.server_name() {
                context.on_server_name(server_name)?;
                self.emitted_server_name = true;
            }
        }
        if !self.emitted_application_protocol {
            if let Some(application_protocol) = self.application_protocol() {
                let application_protocol = Bytes::copy_from_slice(application_protocol);
                context.on_application_protocol(application_protocol)?;
                self.emitted_application_protocol = true;
            }
        }

        Ok(())
    }
}

impl crypto::CryptoSuite for Session {
    type HandshakeKey = PacketKeys;
    type HandshakeHeaderKey = HeaderProtectionKeys;
    type InitialKey = s2n_quic_crypto::initial::InitialKey;
    type InitialHeaderKey = s2n_quic_crypto::initial::InitialHeaderKey;
    type OneRttKey = OneRttKey;
    type OneRttHeaderKey = HeaderProtectionKeys;
    type ZeroRttKey = PacketKey;
    type ZeroRttHeaderKey = HeaderProtectionKey;
    type RetryKey = s2n_quic_crypto::retry::RetryKey;
}

impl tls::Session for Session {
    fn poll<C: tls::Context<Self>>(
        &mut self,
        context: &mut C,
    ) -> Poll<Result<(), transport::Error>> {
        let result = self.poll_impl(context);
        // attempt to emit server_name and application_protocol events prior to possibly
        // returning with an error
        self.emit_events(context)?;
        result
    }
}

#[derive(Clone, Copy, Debug, PartialEq, PartialOrd)]
enum HandshakePhase {
    Initial,
    Handshake,
    Application,
}

impl HandshakePhase {
    fn transition(&mut self) {
        *self = match self {
            Self::Initial => Self::Handshake,
            _ => Self::Application,
        };
    }
}

impl Default for HandshakePhase {
    fn default() -> Self {
        Self::Initial
    }
}


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

use crate::{
    run::Config as Sim,
    stats::{Filter, Query},
    Result,
};
use anyhow::anyhow;
use serde::Deserialize;
use std::{
    collections::BTreeMap,
    fs, io,
    path::{Path, PathBuf},
    process::Command,
    str::FromStr,
};
use structopt::StructOpt;

static INDEX: &str = include_str!("./batch.html");

#[derive(Debug, StructOpt)]
pub struct Batch {
    plans: Vec<Plan>,

    #[structopt(short, long, default_value = "target/s2n-quic-sim")]
    out: PathBuf,

    #[structopt(long)]
    skip_run: bool,
}

impl Batch {
    pub fn run(&self) -> Result {
        let out = &self.out;
        let command = std::env::args().next().unwrap();

        fs::create_dir_all(out)?;
        fs::write(out.join("index.html"), INDEX)?;

        let mut reports = vec![];

        for plan in &self.plans {
            plan.run(out, &command, self.skip_run, &mut reports)?;
        }

        for (_title, report) in reports.iter_mut() {
            *report = report.strip_prefix(out).unwrap().to_owned();
        }

        serde_json::to_writer(fs::File::create(out.join("reports.json"))?, &reports)?;

        Ok(())
    }
}

#[derive(Clone, Debug, Deserialize)]
struct Plan {
    #[serde(default)]
    name: Option<String>,

    sim: Sim,

    report: BTreeMap<String, Report>,
}

impl Plan {
    fn run(
        &self,
        out: &Path,
        command: &str,
        skip_run: bool,
        reports: &mut Vec<(String, PathBuf)>,
    ) -> Result {
        let name = self.name.as_ref().unwrap();
        eprintln!("     Running {name}");
        let out = out.join(name);
        fs::create_dir_all(&out)?;

        let db = out.join("db.proto");

        if !skip_run || !db.exists() {
            let status = Command::new(command)
                .arg("run")
                .arg("--progress")
                .args(self.sim.args())
                .stdout(fs::File::create(&db)?)
                .status()?;

            if !status.success() {
                return Err(anyhow!("run did not succeed"));
            }
        }

        for (report_name, report) in self.report.iter() {
            let mut res = report.run(&out, command, &db, report_name)?;
            res.0 = format!("{} - {}", name, res.0);
            reports.push(res);
        }

        Ok(())
    }
}

impl FromStr for Plan {
    type Err = io::Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let file = fs::read_to_string(s)?;
        let mut plan: Self =
            toml::from_str(&file).map_err(|err| io::Error::new(io::ErrorKind::Other, err))?;

        if plan.name.is_none() {
            plan.name = Some(
                Path::new(s.trim_end_matches(".toml"))
                    .file_name()
                    .and_then(|v| v.to_str())
                    .unwrap_or(s)
                    .to_owned(),
            );
        }

        Ok(plan)
    }
}

#[derive(Clone, Debug, Deserialize)]
struct Report {
    #[serde(default)]
    title: Option<String>,
    #[serde(default)]
    filters: Vec<Filter>,
    x: Query,
    y: Query,
}

impl Report {
    fn run(&self, out: &Path, command: &str, db: &Path, name: &str) -> Result<(String, PathBuf)> {
        let title = self.title.as_deref().unwrap_or(name).to_owned();
        let output = out.join(format!("{name}.json"));

        let mut cmd = Command::new(command);

        cmd.arg("report")
            .arg(db)
            .arg("--x")
            .arg(self.x.to_string())
            .arg("--y")
            .arg(self.y.to_string())
            .arg("--title")
            .arg(&title)
            .stdout(fs::File::create(&output)?);

        for filter in &self.filters {
            cmd.arg("--filter").arg(filter.to_string());
        }

        let status = cmd.status()?;

        if !status.success() {
            return Err(anyhow!("{} report did not succeed", name));
        }

        Ok((title, output))
    }
}


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

pub use anyhow::Error;
use structopt::StructOpt;

pub type Result<T = (), E = Error> = core::result::Result<T, E>;

mod batch;
mod query;
mod report;
mod run;
mod stats;

#[derive(Debug, StructOpt)]
enum Args {
    Query(query::Query),
    Run(Box<run::Run>),
    Report(report::Report),
    Batch(batch::Batch),
}

fn main() -> Result {
    let format = tracing_subscriber::fmt::format()
        .with_level(false) // don't include levels in formatted output
        .with_timer(Uptime)
        .with_ansi(false)
        .compact(); // Use a less verbose output format.

    tracing_subscriber::fmt()
        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
        .event_format(format)
        .init();

    match Args::from_args() {
        Args::Query(args) => args.run(),
        Args::Run(args) => args.run(),
        Args::Report(args) => args.run(),
        Args::Batch(args) => args.run(),
    }
}

struct Uptime;

impl tracing_subscriber::fmt::time::FormatTime for Uptime {
    fn format_time(&self, w: &mut tracing_subscriber::fmt::format::Writer<'_>) -> std::fmt::Result {
        write!(w, "{}", s2n_quic::provider::io::testing::now())
    }
}


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

use crate::{
    stats::{self, Connection, Parameters, Stats},
    Result,
};
use anyhow::anyhow;
use std::{collections::HashMap, fs, io};
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
pub struct Query {
    #[structopt(long)]
    header: Option<Option<bool>>,

    #[structopt(long)]
    clients: Option<Option<bool>>,

    #[structopt(long)]
    servers: Option<Option<bool>>,

    #[structopt(long, short)]
    filter: Vec<stats::Filter>,

    #[structopt(long, short, possible_values = &*stats::QUERY_NAMES)]
    query: Vec<stats::Query>,

    #[structopt(long)]
    with_seed: bool,

    input: Option<String>,
}

impl Query {
    pub fn run(&self) -> Result {
        let input: Box<dyn io::Read> = if self.input.as_ref().map_or(true, |v| v == "-") {
            let reader = io::stdin();
            Box::new(reader)
        } else {
            let reader = fs::File::open(self.input.as_ref().unwrap())?;
            let reader = io::BufReader::new(reader);
            Box::new(reader)
        };

        if self.query.is_empty() {
            return Err(anyhow!("must specify at least one query"));
        }

        let with_seed = self.with_seed;

        let queries = &self.query;
        let filters = &self.filter;

        let header = self.header.map_or(true, |v| v.unwrap_or(true));
        if header {
            let seed = if with_seed {
                Some(("seed", "id"))
            } else {
                None
            };
            let queries = queries.iter().map(|q| q.name);
            if emit(seed, queries).is_err() {
                return Ok(());
            }
        }

        let mut values = vec![];

        let mut q = |p: &Parameters, conn: &Connection, connections: &[Connection]| -> Result {
            for filter in filters {
                if !filter.apply(p, conn, connections) {
                    return Ok(());
                }
            }

            values.clear();

            for query in queries {
                if let Some(value) = query.apply(p, conn, connections) {
                    values.push(value);
                } else {
                    return Ok(());
                }
            }

            let seed = if with_seed {
                Some((conn.seed, conn.id()))
            } else {
                None
            };

            emit(seed, values.iter())?;

            Ok(())
        };

        let clients = self.clients.map_or(true, |v| v.unwrap_or(true));
        let servers = self.servers.map_or(true, |v| v.unwrap_or(true));

        let reader = crate::stats::Stats::reader(input);

        let mut acc: HashMap<u64, Vec<Connection>> = HashMap::new();

        for stat in reader {
            match stat? {
                Stats::Setup(_) => {
                    // unused
                }
                Stats::Parameters(p) => {
                    if let Some(connections) = acc.remove(&p.seed) {
                        for conn in &connections {
                            if q(&p, conn, &connections).is_err() {
                                return Ok(());
                            }
                        }
                    }
                }
                Stats::Connection(c) => {
                    if (clients && c.client_id.is_some()) || (servers && c.server_id.is_some()) {
                        acc.entry(c.seed).or_default().push(c);
                    }
                }
            }
        }

        Ok(())
    }
}

fn emit<
    Seed: core::fmt::Display,
    Id: core::fmt::Display,
    I: Iterator<Item = V>,
    V: core::fmt::Display,
>(
    seed: Option<(Seed, Id)>,
    i: I,
) -> io::Result<()> {
    use io::Write;
    let stdout = io::stdout();
    let mut o = stdout.lock();

    let mut has_written = false;

    if let Some((seed, id)) = seed {
        write!(o, "{seed}\t{id}")?;
        has_written = true;
    }

    for value in i {
        if has_written {
            write!(o, "\t{value}")?;
        } else {
            write!(o, "{value}")?;
        }
        has_written = true;
    }
    writeln!(o)?;
    Ok(())
}


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

use crate::{
    stats::{self, Connection, Filter, Parameters, Query, Stats, QUERY_NAMES},
    Result,
};
use serde_json::json;
use std::{
    collections::{BTreeSet, HashMap},
    fs, io,
    path::PathBuf,
};
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
pub struct Report {
    #[structopt(long, short)]
    filter: Vec<Filter>,

    #[structopt(long, short, possible_values = &*QUERY_NAMES)]
    x: Query,

    #[structopt(long, default_value = "100")]
    x_width: u32,

    #[structopt(long, short, possible_values = &*QUERY_NAMES)]
    y: Query,

    #[structopt(long, default_value = "100")]
    y_width: u32,

    #[structopt(long)]
    title: Option<String>,

    #[structopt(long, default_value = "blues")]
    palette: String,

    input: PathBuf,
}

impl Report {
    pub fn run(&self) -> Result {
        let reader = fs::File::open(&self.input)?;
        let reader = io::BufReader::new(reader);
        let reader = Stats::reader(reader);

        let filters = &self.filter;

        let mut seed_ids: HashMap<u64, usize> = HashMap::new();
        let mut seed_hist = vec![];
        let mut acc: HashMap<usize, Vec<Connection>> = HashMap::new();
        let mut values = vec![];
        let mut x_bounds = Bounds::default();
        let mut y_bounds = Bounds::default();

        let mut seed_id = |seed: u64| {
            use std::collections::hash_map::Entry;
            let next_id = seed_ids.len();
            match seed_ids.entry(seed) {
                Entry::Occupied(entry) => *entry.get(),
                Entry::Vacant(entry) => {
                    entry.insert(next_id);

                    seed_hist.push((seed, vec![], vec![]));

                    next_id
                }
            }
        };

        let mut q =
            |seed_id: usize, p: &Parameters, conn: &Connection, connections: &[Connection]| {
                for filter in filters {
                    if !filter.apply(p, conn, connections) {
                        return;
                    }
                }

                let x = self.x.apply(p, conn, connections);
                let y = self.y.apply(p, conn, connections);

                if let (Some(x), Some(y)) = (x, y) {
                    x_bounds.push(x);
                    y_bounds.push(y);
                    values.push((seed_id, conn.id(), x, y));
                }
            };

        let mut args = vec![];

        for stat in reader {
            match stat? {
                Stats::Setup(s) => {
                    args = s.args;
                }
                Stats::Parameters(p) => {
                    let id = seed_id(p.seed);
                    if let Some(connections) = acc.remove(&id) {
                        for conn in &connections {
                            q(id, &p, conn, &connections);
                        }
                    }
                }
                Stats::Connection(c) => {
                    let id = seed_id(c.seed);
                    acc.entry(id).or_default().push(c);
                }
            }
        }

        let x_width = self.x_width;
        let y_width = self.y_width;
        let mut hist = vec![(0u32, 0u32, BTreeSet::new()); x_width as usize * y_width as usize];

        let x_bin = x_bounds.bin(x_width as usize);
        let y_bin = y_bounds.bin(y_width as usize);

        let mut max_depth = 0;
        for (seed, _id, x, y) in values.iter().copied() {
            let x = x_bin.bin(x);
            let y = y_bin.bin(y);

            let idx = x + y * self.x_width as usize;
            let entry = &mut hist[idx];
            let value = entry.1 + 1;

            entry.0 = idx as _;
            entry.1 = value;
            entry.2.insert(seed);
            max_depth = max_depth.max(value);

            let entry = &mut seed_hist[seed];
            if entry.1.is_empty() {
                entry.1 = vec![(0u32, 0u32); x_width as usize];
            }
            if entry.2.is_empty() {
                entry.2 = vec![(0u32, 0u32); y_width as usize];
            }

            // increment the x hist
            entry.1[x].0 = x as _;
            entry.1[x].1 += 1;
            // increment the y hist
            entry.2[y].0 = y as _;
            entry.2[y].1 += 1;
        }

        //return Ok(());

        // remove any unused entries
        hist.retain(|(_idx, count, _seeds)| *count > 0);

        for (_seed, x_hist, y_hist) in seed_hist.iter_mut() {
            x_hist.retain(|(_idx, count)| *count > 0);
            y_hist.retain(|(_idx, count)| *count > 0);
        }

        let title = self
            .title
            .clone()
            .unwrap_or_else(|| format!("{} vs {}", self.x, self.y));

        let total_height = 1200;
        let total_width = 1200;
        let size = 600;
        let hist_width = size / self.x_width;
        let hist_height = size / self.y_width;

        let group_width = 40;
        let group_height = 15;
        let group_padding = 60;
        let group_offset = size + group_padding;

        let seed_hist_size = size - group_padding;
        let groups_per_line = seed_hist_size / group_width;

        let command = format!(
            "{:?} + (isNumber({s}) ? (' --seed ' + {s}) : '')",
            args.join(" "),
            s = "sig$seed_hist.seed"
        );

        let vega = json!({
            "$schema": "https://vega.github.io/schema/vega/v5.json",
            "width": total_width,
            "height": total_height,
            "padding": 5,
            "background": "white",

            "title": {
                "text": title,
                "anchor": "middle",
                "fontSize": 16,
                "frame": "group",
                "offset": 4
            },

            "signals": [
                {
                    "name": "sig$group",
                    "value": [],
                    "on": [{ "events": "@hist:click", "update": "datum.group" }]
                },
                {
                    "name": "sig$seed",
                    "on": [
                        { "events": "@hist:click", "update": "0" },
                        { "events": "@seed_id:click", "update": "datum.index" }
                    ]
                },
                {
                    "name": "sig$seed_hist",
                    "update": "isNumber(sig$seed) ? data('data$seed_hist')[sig$group[sig$seed]] : {}"
                },
                {
                    "name": "sig$command",
                    "update": command,
                }
            ],

            "scales": [
                {
                    "name": "scale$hist_x",
                    "type": "linear",
                    "domain": [0, self.x_width + 1],
                    "range": [0, size]
                },
                {
                    "name": "scale$axis_x",
                    "type": "linear",
                    "domain": x_bounds.domain(self.x.ty),
                    "range": [0, size],
                    "nice": true,
                    "zero": false,
                },
                {
                    "name": "scale$hist_y",
                    "type": "linear",
                    "domain": [0, self.y_width + 1],
                    "range": [0, size]
                },
                {
                    "name": "scale$axis_y",
                    "type": "linear",
                    "domain": y_bounds.domain(self.y.ty),
                    "range": [0, size],
                    "reverse": true,
                    "nice": true,
                    "zero": false,
                },
                {
                    "name": "scale$seed_hist_x",
                    "type": "linear",
                    "domain": [0, { "signal": "sig$seed_hist_x_extent[1]" }],
                    "range": [0, seed_hist_size]
                },
                {
                    "name": "scale$seed_axis_x",
                    "type": "linear",
                    "domain": { "signal": "sig$seed_hist_x_extent" },
                    "range": [group_offset, total_height],
                    "nice": true,
                    "zero": false,
                },
                {
                    "name": "scale$seed_hist_y",
                    "type": "linear",
                    "domain": [0, { "signal": "sig$seed_hist_y_extent[1]" }],
                    "range": [0, seed_hist_size]
                },
                {
                    "name": "scale$seed_axis_y",
                    "type": "linear",
                    "domain": { "signal": "sig$seed_hist_y_extent" },
                    "range": [group_offset, total_height],
                    "nice": true,
                    "zero": false,
                },
                {
                    "name": "scale$color",
                    "type": "linear",
                    "range": { "scheme": self.palette },
                    "domain": [0, max_depth]
                }
            ],

            "legends": [
                {
                    "orient": "none",
                    "fill": "scale$color",
                    "type": "gradient",
                    "direction": "horizontal",
                    "gradientLength": size,
                    "legendY": size + 20,
                    "legendX": 10
                }
            ],

            "axes": [
                {
                    "orient": "top",
                    "scale": "scale$axis_x",
                    "domain": false,
                    "title": self.x.to_string(),
                    "format": self.x.ty.format(x_bounds.domain(self.x.ty)),
                    "formatType": if self.x.ty.is_duration() {
                        "time"
                    } else {
                        "number"
                    },
                },
                {
                    "orient": "bottom",
                    "scale": "scale$axis_x",
                    "domain": false,
                    "title": self.x.to_string(),
                    "format": self.x.ty.format(x_bounds.domain(self.x.ty)),
                    "formatType": if self.x.ty.is_duration() {
                        "time"
                    } else {
                        "number"
                    },
                },
                {
                    "orient": "left",
                    "scale": "scale$seed_axis_x",
                    "domain": false,
                    "title": "Count",
                },
                {
                    "orient": "left",
                    "scale": "scale$axis_y",
                    "domain": false,
                    "title": self.y.to_string(),
                    "format": self.y.ty.format(y_bounds.domain(self.y.ty)),
                    "formatType": if self.y.ty.is_duration() {
                        "time"
                    } else {
                        "number"
                    },
                },
                {
                    "orient": "top",
                    "scale": "scale$seed_axis_y",
                    "domain": false,
                    "title": "Count",
                },
                {
                    "orient": "right",
                    "scale": "scale$axis_y",
                    "domain": false,
                    "title": self.y.to_string(),
                    "format": self.y.ty.format(y_bounds.domain(self.y.ty)),
                    "formatType": if self.y.ty.is_duration() {
                        "time"
                    } else {
                        "number"
                    },
                }
            ],

            "marks": [
                {
                    "name": "hist",
                    "type": "rect",
                    "from": { "data": "data$hist" },
                    "encode": {
                        "enter": {
                            "width": { "value": hist_width },
                            "height": { "value": hist_height },
                        },
                        "update": {
                            "x": { "scale": "scale$hist_x", "signal": "datum.x" },
                            "y": { "scale": "scale$hist_y", "signal": "datum.y" },
                            "fill": { "scale": "scale$color", "field": "count" },
                            "tooltip": { "signal": "datum.count" },
                        },
                    }
                },
                {
                    "name": "seed_hist_x",
                    "type": "rect",
                    "from": { "data": "data$seed_hist_x" },
                    "encode": {
                        "enter": {
                            "y": { "value": group_offset },
                            "width": { "value": hist_width },
                        },
                        "update": {
                            "x": { "scale": "scale$hist_x", "signal": "datum.x" },
                            "height": { "scale": "scale$seed_hist_x", "field": "count" },
                            "fill": { "scale": "scale$color", "field": "count" },
                            "tooltip": { "signal": "datum.count" },
                        },
                    }
                },
                {
                    "name": "seed_hist_y",
                    "type": "rect",
                    "from": { "data": "data$seed_hist_y" },
                    "encode": {
                        "enter": {
                            "x": { "value": group_offset },
                            "height": { "value": hist_height },
                        },
                        "update": {
                            "y": { "scale": "scale$hist_y", "signal": "datum.y" },
                            "width": { "scale": "scale$seed_hist_y", "field": "count" },
                            "fill": { "scale": "scale$color", "field": "count" },
                            "tooltip": { "signal": "datum.count" },
                        },
                    }
                },
                {
                    "type": "text",
                    "from": { "data": "data$group" },
                    "encode": {
                        "enter": {
                            "width": { "value": group_width },
                            "height": { "value": group_height },
                        },
                        "update": {
                            "x": { "signal": "datum.x" },
                            "y": { "signal": "datum.y" },
                            "text": { "signal": "datum.index + 1" },
                        }
                    }
                },
                {
                    "name": "seed_id",
                    "type": "rect",
                    "from": { "data": "data$group" },
                    "encode": {
                        "enter": {
                            "width": { "value": group_width },
                            "height": { "value": group_height },
                            "fill": { "value": "#125ca4" },
                        },
                        "update": {
                            "x": { "signal": "datum.x - 5" },
                            "y": { "signal": "datum.y - 12" },
                            "fillOpacity": { "signal": "sig$seed == datum.index ? 0.3 : 0.0" },
                        }
                    }
                },
            ],

            "data": [
                {
                    "name": "data$hist",
                    "values": &hist,
                    "transform": [
                        {
                            "type": "formula",
                            "as": "index",
                            "expr": "datum[0]",
                        },
                        {
                            "type": "formula",
                            "as": "count",
                            "expr": "datum[1]",
                        },
                        {
                            "type": "formula",
                            "as": "group",
                            "expr": "datum[2]",
                        },
                        {
                            "type": "formula",
                            "as": "x",
                            "expr": format!("floor(datum.index % {x_width})"),
                        },
                        {
                            "type": "formula",
                            "as": "y",
                            "expr": format!("{y_width} - floor(datum.index / {y_width})"),
                        },
                    ],
                },
                {
                    "name": "data$seed_hist",
                    "values": &seed_hist,
                    "transform": [
                        {
                            "type": "formula",
                            "as": "seed",
                            "expr": "datum[0]",
                        },
                        {
                            "type": "formula",
                            "as": "x_hist",
                            "expr": "datum[1]",
                        },
                        {
                            "type": "formula",
                            "as": "y_hist",
                            "expr": "datum[2]",
                        },
                    ],
                },
                {
                    "name": "data$seed_hist_x",
                    "values": { "signal": "sig$seed_hist.x_hist" },
                    "transform": [
                        {
                            "type": "formula",
                            "as": "index",
                            "expr": "datum[0]",
                        },
                        {
                            "type": "formula",
                            "as": "count",
                            "expr": "datum[1]",
                        },
                        {
                            "type": "formula",
                            "as": "x",
                            "expr": "datum.index",
                        },
                        {
                            "type": "extent",
                            "field": "count",
                            "signal": "sig$seed_hist_x_extent"
                        },
                    ],
                },
                {
                    "name": "data$seed_hist_y",
                    "values": { "signal": "sig$seed_hist.y_hist" },
                    "transform": [
                        {
                            "type": "formula",
                            "as": "index",
                            "expr": "datum[0]",
                        },
                        {
                            "type": "formula",
                            "as": "count",
                            "expr": "datum[1]",
                        },
                        {
                            "type": "formula",
                            "as": "y",
                            "expr": format!("{y_width} - datum.index"),
                        },
                        {
                            "type": "extent",
                            "field": "count",
                            "signal": "sig$seed_hist_y_extent"
                        },
                    ],
                },
                {
                    "name": "data$group",
                    "transform": [
                        {
                            "type": "sequence",
                            "start": 0,
                            "stop": { "signal": "length(sig$group)" },
                            "as": "index",
                        },
                        {
                            "type": "formula",
                            "as": "x",
                            "expr": format!("{} + floor(datum.index % {}) * {}", group_offset + 20, groups_per_line, group_width),
                        },
                        {
                            "type": "formula",
                            "as": "y",
                            "expr": format!("{} + floor(datum.index / {}) * {}", group_offset + 20, groups_per_line, group_height),
                        },
                    ]
                },
            ],
        });

        println!("{vega}");

        Ok(())
    }
}

#[derive(Clone, Copy, Debug)]
struct Bounds {
    min: f64,
    max: f64,
}

impl Default for Bounds {
    fn default() -> Self {
        Self {
            min: f64::MAX,
            max: f64::MIN,
        }
    }
}

impl Bounds {
    pub fn push(&mut self, value: f64) {
        self.min = self.min.min(value);
        self.max = self.max.max(value);
    }

    pub fn bin(&self, width: usize) -> Bin {
        Bin {
            min: self.min,
            step: (self.max - self.min) / (width - 1) as f64,
        }
    }

    pub fn domain(&self, ty: stats::Type) -> [f64; 2] {
        if ty.is_duration() {
            [self.min * 1000.0, self.max * 1000.0]
        } else {
            [self.min, self.max]
        }
    }
}

#[derive(Clone, Copy, Debug)]
struct Bin {
    min: f64,
    step: f64,
}

impl Bin {
    pub fn bin(&self, value: f64) -> usize {
        ((value - self.min) / self.step) as usize
    }
}


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

use crate::{stats, Result};
use indicatif::{ParallelProgressIterator, ProgressBar};
use rayon::prelude::*;
use s2n_quic::provider::io::testing::{test_seed, Model};
use structopt::StructOpt;

mod config;
pub use config::Config;

mod endpoint;
mod events;

mod range;
use range::CliRange;

#[derive(Debug, StructOpt)]
pub struct Run {
    #[structopt(flatten)]
    config: Config,

    #[structopt(long)]
    seed: Vec<u64>,

    #[structopt(long)]
    progress: bool,
}

impl core::ops::Deref for Run {
    type Target = Config;

    fn deref(&self) -> &Self::Target {
        &self.config
    }
}

impl Run {
    pub fn run(&self) -> Result {
        assert_ne!(self.servers.start, 0);
        assert_ne!(self.clients.start, 0);
        assert_ne!(self.connections.start, 0);

        let test = |seed: u64| {
            let network = Model::default();

            test_seed(network.clone(), seed, |handle| {
                let server_len = self.servers.gen();
                let client_len = self.clients.gen();

                let events = self.gen_network(seed, server_len, client_len, &network);

                let mut servers = vec![];
                for _ in 0..server_len {
                    servers.push(endpoint::server(handle, events.clone())?);
                }

                for _ in 0..client_len {
                    let count = self.connections.gen() as usize;
                    let delay = self.connect_delay;
                    let streams = self.streams;
                    let stream_data = self.stream_data;
                    endpoint::client(
                        handle,
                        events.clone(),
                        &servers,
                        count,
                        delay,
                        streams,
                        stream_data,
                    )?;
                }

                Ok(())
            })
            .unwrap();
        };

        if self.seed.is_empty() {
            events::dump(|stdout| {
                stats::Setup {
                    args: self.config.args(),
                }
                .write(stdout)
            });

            let pb = if self.progress {
                ProgressBar::new(self.iterations as _)
            } else {
                ProgressBar::hidden()
            };

            static MSG: &str =
                "{spinner:.green} [{elapsed_precise}] [{bar:.cyan/blue}] {pos}/{len} ({eta})";

            pb.set_style(
                indicatif::ProgressStyle::default_bar()
                    .template(MSG)
                    .unwrap()
                    .progress_chars("=> "),
            );

            (0..self.iterations)
                .into_par_iter()
                .progress_with(pb)
                .map(|v| if events::is_open() { Some(v) } else { None })
                .while_some()
                .for_each(|_| {
                    use ::rand::prelude::*;
                    let seed = thread_rng().gen();
                    test(seed);
                });
        } else {
            // don't dump events when running specific seeds
            events::close();

            for seed in self.seed.iter().copied() {
                test(seed);
            }
        }

        Ok(())
    }

    fn gen_network(&self, seed: u64, servers: u32, clients: u32, model: &Model) -> events::Events {
        let mut events = stats::Parameters {
            seed,
            servers,
            clients,
            ..Default::default()
        };

        macro_rules! param {
            ($name:ident, $set:ident, gen_duration) => {{
                let value = self.$name.gen_duration();
                model.$set(value);
                events.$name = Some(value.into());
            }};
            ($name:ident, $set:ident, $gen:ident $($tt:tt)*) => {{
                let value = self.$name.$gen();
                model.$set(value);
                events.$name = value $($tt)*;
            }};
        }

        param!(drop_rate, set_drop_rate, gen * 100.0);
        param!(corrupt_rate, set_corrupt_rate, gen * 100.0);
        param!(jitter, set_jitter, gen_duration);
        param!(network_jitter, set_network_jitter, gen_duration);
        param!(delay, set_delay, gen_duration);
        param!(inflight_delay, set_inflight_delay, gen_duration);
        param!(retransmit_rate, set_retransmit_rate, gen * 100.0);
        param!(max_udp_payload, set_max_udp_payload, gen as _);

        macro_rules! zero_param {
            ($name:ident, $set:ident) => {
                if self.$name.end > 0 {
                    let value = self.$name.gen();
                    model.$set(value);
                    events.$name = value;
                } else {
                    events.$name = u64::MAX;
                }
            };
        }

        zero_param!(transmit_rate, set_transmit_rate);
        zero_param!(max_inflight, set_max_inflight);
        zero_param!(inflight_delay_threshold, set_inflight_delay_threshold);

        events.into()
    }
}


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

use core::fmt;
use once_cell::sync::Lazy;
use prost::Message;
use s2n_quic::provider::event::{events, Timestamp};
use std::{
    cell::RefCell,
    convert::TryInto,
    io::{self, Read},
    str::FromStr,
};

thread_local! {
    static BUFFER: RefCell<Vec<u8>> = Default::default();
}

#[allow(clippy::large_enum_variant)] // the small variants are only used once so it's not a big deal
#[derive(Debug)]
pub enum Stats {
    Setup(Setup),
    Parameters(Parameters),
    Connection(Connection),
}

impl Stats {
    pub fn reader<R: io::Read>(read: R) -> impl Iterator<Item = io::Result<Self>> {
        Reader {
            read: io::BufReader::new(read),
        }
    }

    fn write<W: io::Write, M: Message>(mut w: W, tag: u8, msg: &M) -> io::Result<()> {
        BUFFER.with(|cell| {
            let mut buffer = cell.borrow_mut();
            let buffer = &mut *buffer;
            buffer.clear();

            buffer.push(tag);
            let len = msg.encoded_len();
            let len: u16 = len.try_into().unwrap();
            buffer.extend_from_slice(&len.to_le_bytes());
            msg.encode(buffer).unwrap();

            w.write_all(buffer)?;

            Ok(())
        })
    }
}

pub struct Reader<R: io::Read> {
    read: io::BufReader<R>,
}

impl<R: io::Read> Reader<R> {
    fn read(&mut self) -> io::Result<Stats> {
        let mut prefix = [0u8, 0, 0];
        self.read.read_exact(&mut prefix)?;

        let id = prefix[0];
        let len = u16::from_le_bytes([prefix[1], prefix[2]]) as usize;

        BUFFER.with(|cell| {
            let mut buffer = cell.borrow_mut();
            let buffer = &mut *buffer;
            buffer.resize(len, 0);

            self.read.read_exact(buffer)?;

            let buffer = io::Cursor::new(buffer);

            match id {
                0 => {
                    let msg = Parameters::decode(buffer)?;
                    Ok(msg.into())
                }
                1 => {
                    let msg = Connection::decode(buffer)?;
                    Ok(msg.into())
                }
                2 => {
                    let msg = Setup::decode(buffer)?;
                    Ok(msg.into())
                }
                _ => Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "invalid stat tag",
                )),
            }
        })
    }
}

impl<R: io::Read> Iterator for Reader<R> {
    type Item = io::Result<Stats>;

    fn next(&mut self) -> Option<Self::Item> {
        match self.read() {
            Ok(value) => Some(Ok(value)),
            Err(err) if err.kind() == io::ErrorKind::UnexpectedEof => None,
            Err(err) => Some(Err(err)),
        }
    }
}

#[derive(Clone, Message)]
pub struct Setup {
    #[prost(string, repeated, tag = "1")]
    pub args: Vec<String>,
}

impl From<Setup> for Stats {
    fn from(s: Setup) -> Self {
        Self::Setup(s)
    }
}

impl Setup {
    pub fn write<W: io::Write>(&self, w: W) -> io::Result<()> {
        Stats::write(w, 2, self)
    }
}

#[derive(Clone, Copy, Message)]
pub struct Parameters {
    #[prost(uint64, tag = "1")]
    pub seed: u64,
    #[prost(double, tag = "2")]
    pub drop_rate: f64,
    #[prost(double, tag = "3")]
    pub corrupt_rate: f64,
    #[prost(message, tag = "4")]
    pub jitter: Option<Duration>,
    #[prost(message, tag = "5")]
    pub network_jitter: Option<Duration>,
    #[prost(message, tag = "6")]
    pub delay: Option<Duration>,
    #[prost(double, tag = "7")]
    pub retransmit_rate: f64,
    #[prost(uint32, tag = "8")]
    pub max_udp_payload: u32,
    #[prost(uint64, tag = "9")]
    pub transmit_rate: u64,
    #[prost(uint64, tag = "10")]
    pub max_inflight: u64,
    #[prost(uint32, tag = "11")]
    pub servers: u32,
    #[prost(uint32, tag = "12")]
    pub clients: u32,
    #[prost(message, tag = "13")]
    pub end_time: Option<Duration>,
    #[prost(message, tag = "14")]
    pub inflight_delay: Option<Duration>,
    #[prost(uint64, tag = "15")]
    pub inflight_delay_threshold: u64,
}

impl From<Parameters> for Stats {
    fn from(p: Parameters) -> Self {
        Self::Parameters(p)
    }
}

impl Parameters {
    pub fn write<W: io::Write>(&self, w: W) -> io::Result<()> {
        Stats::write(w, 0, self)
    }
}

#[derive(Clone, Copy, Message)]
pub struct Connection {
    #[prost(uint64, optional, tag = "1")]
    pub client_id: Option<u64>,
    #[prost(uint64, optional, tag = "2")]
    pub server_id: Option<u64>,
    #[prost(uint64, tag = "3")]
    pub seed: u64,
    #[prost(message, tag = "4")]
    pub start_time: Option<Duration>,
    #[prost(message, tag = "5")]
    pub end_time: Option<Duration>,
    #[prost(message, tag = "6")]
    pub handshake: Option<Handshake>,
    #[prost(uint64, optional, tag = "7")]
    pub transport_error: Option<u64>,
    #[prost(uint64, optional, tag = "8")]
    pub application_error: Option<u64>,
    #[prost(bool, tag = "9")]
    pub idle_timer_error: bool,
    #[prost(bool, tag = "10")]
    pub handshake_duration_exceeded_error: bool,
    #[prost(bool, tag = "11")]
    pub unspecified_error: bool,
    #[prost(message, tag = "12")]
    pub tx: Option<Counts>,
    #[prost(message, tag = "13")]
    pub rx: Option<Counts>,
    #[prost(message, tag = "14")]
    pub loss: Option<Counts>,
    #[prost(uint64, tag = "15")]
    pub congestion: u64,
    #[prost(uint64, tag = "16")]
    pub max_cwin: u64,
    #[prost(uint64, tag = "17")]
    pub max_bytes_in_flight: u64,
    #[prost(message, tag = "18")]
    pub max_rtt: Option<Duration>,
    #[prost(message, tag = "19")]
    pub min_rtt: Option<Duration>,
    #[prost(message, tag = "20")]
    pub smoothed_rtt: Option<Duration>,
}

impl From<Connection> for Stats {
    fn from(c: Connection) -> Self {
        Self::Connection(c)
    }
}

impl Connection {
    pub fn write<W: io::Write>(&self, w: W) -> io::Result<()> {
        Stats::write(w, 1, self)
    }

    pub fn id(&self) -> u64 {
        self.client_id.or(self.server_id).expect("missing id")
    }

    pub fn is_success(&self) -> bool {
        !self.is_error()
    }

    pub fn is_error(&self) -> bool {
        self.transport_error.is_some()
            || self.application_error.is_some()
            || self.idle_timer_error
            || self.handshake_duration_exceeded_error
            || self.unspecified_error
    }

    pub fn duration(&self) -> Option<core::time::Duration> {
        self.end_time
            .unwrap_or_default()
            .as_duration()
            .checked_sub(self.start_time.unwrap_or_default().as_duration())
    }
}

#[derive(Clone, Copy, Message, PartialEq, Eq)]
pub struct Counts {
    #[prost(uint64, tag = "1")]
    pub initial: u64,
    #[prost(uint64, tag = "2")]
    pub handshake: u64,
    #[prost(uint64, tag = "3")]
    pub retry: u64,
    #[prost(uint64, tag = "4")]
    pub one_rtt: u64,
    #[prost(uint64, tag = "5")]
    pub stream_progress: u64,
    #[prost(uint64, tag = "6")]
    pub stream_data_blocked: u64,
    #[prost(uint64, tag = "7")]
    pub data_blocked: u64,
    #[prost(message, tag = "8")]
    pub stream_progress_start: Option<Duration>,
    #[prost(message, tag = "9")]
    pub stream_progress_end: Option<Duration>,
}

impl Counts {
    #[inline]
    pub fn inc_packet(&mut self, header: &events::PacketHeader) {
        match header {
            events::PacketHeader::Initial { .. } => self.initial += 1,
            events::PacketHeader::Handshake { .. } => self.handshake += 1,
            events::PacketHeader::OneRtt { .. } => self.one_rtt += 1,
            _ => {}
        }
    }

    #[inline]
    pub fn inc_frame(&mut self, frame: &events::Frame) {
        use events::Frame::*;
        match frame {
            DataBlocked { .. } => {
                self.data_blocked += 1;
            }
            StreamDataBlocked { .. } => {
                self.stream_data_blocked += 1;
            }
            _ => {}
        }
    }

    #[inline]
    pub fn packets(&self) -> u64 {
        self.initial + self.handshake + self.retry + self.one_rtt
    }

    #[inline]
    pub fn stream_progress(&mut self, now: Timestamp, bytes: usize) {
        self.stream_progress += bytes as u64;
        self.stream_progress_end = Some(now.duration_since_start().into());
        if self.stream_progress_start.is_none() {
            self.stream_progress_start = Some(now.duration_since_start().into());
        }
    }

    #[inline]
    pub fn stream_throughput(&self) -> Option<f64> {
        let bytes = self.stream_progress as f64;
        let start = self.stream_progress_start?.as_duration();
        let end = self.stream_progress_end?.as_duration();
        let duration = end - start;
        Some(bytes / duration.as_secs_f64())
    }
}

#[derive(Clone, Copy, Message, PartialEq, Eq)]
pub struct Handshake {
    #[prost(message, tag = "1")]
    pub complete: Option<Duration>,
    #[prost(message, tag = "2")]
    pub confirmed: Option<Duration>,
}

#[derive(Clone, Copy, Message, PartialEq, Eq)]
pub struct Duration {
    #[prost(uint64, tag = "1")]
    pub secs: u64,
    #[prost(uint32, tag = "2")]
    pub nanos: u32,
}

impl Duration {
    pub fn as_duration(self) -> core::time::Duration {
        self.into()
    }
}

impl From<core::time::Duration> for Duration {
    fn from(value: core::time::Duration) -> Self {
        Self {
            secs: value.as_secs(),
            nanos: value.subsec_nanos(),
        }
    }
}

impl From<Duration> for core::time::Duration {
    fn from(value: Duration) -> Self {
        Self::new(value.secs, value.nanos)
    }
}

type Q = fn(&Parameters, &Connection, &[Connection]) -> Option<f64>;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Type {
    Integer,
    Percent,
    Duration,
    Throughput,
    Bool,
}

impl Type {
    pub fn format(&self, [_min, max]: [f64; 2]) -> &'static str {
        match self {
            Self::Integer => "~s",
            Self::Percent => "~%",
            Self::Duration if max > 2000.0 => "%M:%S",
            Self::Duration => "%Qms",
            Self::Throughput => "~s",
            Self::Bool => "c",
        }
    }

    pub fn is_duration(&self) -> bool {
        matches!(self, Self::Duration)
    }

    pub fn parse(&self, value: &str) -> anyhow::Result<f64> {
        match self {
            Self::Integer => {
                let value: u64 = value.parse()?;
                Ok(value as _)
            }
            Self::Percent => {
                let (value, mul) = if let Some(value) = value.strip_suffix('%') {
                    (value, 100.0)
                } else {
                    (value, 1.0)
                };

                if let Ok(value) = value.parse::<u64>() {
                    let value = value as f64;
                    Ok(value * mul)
                } else {
                    let value = value.parse::<f64>()?;
                    Ok(value * mul)
                }
            }
            Self::Duration => {
                if let Ok(v) = value.parse::<humantime::Duration>() {
                    Ok(v.as_secs_f64())
                } else {
                    Ok(value.parse()?)
                }
            }
            Self::Throughput => {
                todo!()
            }
            Self::Bool => match value {
                "true" | "TRUE" | "1" => Ok(1.0),
                "false" | "FALSE" | "0" => Ok(0.0),
                _ => Err(anyhow::anyhow!("invalid bool: {:?}", value)),
            },
        }
    }
}

use Type::{Bool as B, Duration as T, Integer as I, Percent as P, Throughput as Tpt};

static QUERIES: &[(&str, Type, Q)] = &[
    ("conn.duration", T, |_params, conn, _conns| {
        let duration = conn.duration()?;
        Some(duration.as_secs_f64())
    }),
    ("conn.handshake.confirmed", T, |_params, conn, _conns| {
        let duration = conn
            .handshake
            .unwrap_or_default()
            .confirmed
            .unwrap_or_default()
            .as_duration()
            .checked_sub(conn.start_time.unwrap_or_default().as_duration())?;
        Some(duration.as_secs_f64())
    }),
    ("conn.handshake.complete", T, |_params, conn, _conns| {
        let duration = conn
            .handshake
            .unwrap_or_default()
            .complete
            .unwrap_or_default()
            .as_duration()
            .checked_sub(conn.start_time.unwrap_or_default().as_duration())?;
        Some(duration.as_secs_f64())
    }),
    ("conn.id", B, |_params, conn, _conns| {
        conn.client_id.or(conn.server_id).map(|v| v as f64)
    }),
    ("conn.client", B, |_params, conn, _conns| {
        Some(if conn.client_id.is_some() { 1.0 } else { 0.0 })
    }),
    ("conn.server", B, |_params, conn, _conns| {
        Some(if conn.server_id.is_some() { 1.0 } else { 0.0 })
    }),
    ("conn.success", B, |_params, conn, _conns| {
        Some(if conn.is_success() { 1.0 } else { 0.0 })
    }),
    ("conn.error", B, |_params, conn, _conns| {
        Some(if conn.is_error() { 1.0 } else { 0.0 })
    }),
    ("conn.congestion", I, |_params, conn, _conns| {
        Some(conn.congestion as _)
    }),
    ("conn.max-cwin", I, |_params, conn, _conns| {
        Some(conn.max_cwin as _)
    }),
    ("conn.max-bytes-in-flight", I, |_params, conn, _conns| {
        Some(conn.max_bytes_in_flight as _)
    }),
    ("conn.min-rtt", T, |_params, conn, _conns| {
        Some(conn.min_rtt?.as_duration().as_secs_f64())
    }),
    ("conn.max-rtt", T, |_params, conn, _conns| {
        Some(conn.max_rtt?.as_duration().as_secs_f64())
    }),
    ("conn.smoothed-rtt", T, |_params, conn, _conns| {
        Some(conn.smoothed_rtt?.as_duration().as_secs_f64())
    }),
    ("conn.rtt-spread", T, |_params, conn, _conns| {
        let min = conn.min_rtt?.as_duration();
        let max = conn.max_rtt?.as_duration();
        let variance = max - min;
        Some(variance.as_secs_f64())
    }),
    ("conn.tx.packets", I, |_params, conn, _conns| {
        Some(conn.tx.unwrap_or_default().packets() as _)
    }),
    ("conn.tx.initial", I, |_params, conn, _conns| {
        Some(conn.tx.unwrap_or_default().initial as _)
    }),
    ("conn.tx.handshake", I, |_params, conn, _conns| {
        Some(conn.tx.unwrap_or_default().handshake as _)
    }),
    ("conn.tx.one-rtt", I, |_params, conn, _conns| {
        Some(conn.tx.unwrap_or_default().one_rtt as _)
    }),
    ("conn.tx.data-blocked", I, |_params, conn, _conns| {
        Some(conn.tx.unwrap_or_default().data_blocked as _)
    }),
    ("conn.tx.stream-data-blocked", I, |_params, conn, _conns| {
        Some(conn.tx.unwrap_or_default().stream_data_blocked as _)
    }),
    ("conn.tx.stream-progress", I, |_params, conn, _conns| {
        Some(conn.tx.unwrap_or_default().stream_progress as _)
    }),
    ("conn.tx.stream-throughput", Tpt, |_params, conn, _conns| {
        conn.tx?.stream_throughput()
    }),
    ("conn.rx.packets", I, |_params, conn, _conns| {
        Some(conn.rx.unwrap_or_default().packets() as _)
    }),
    ("conn.rx.initial", I, |_params, conn, _conns| {
        Some(conn.rx.unwrap_or_default().initial as _)
    }),
    ("conn.rx.handshake", I, |_params, conn, _conns| {
        Some(conn.rx.unwrap_or_default().handshake as _)
    }),
    ("conn.rx.one-rtt", I, |_params, conn, _conns| {
        Some(conn.rx.unwrap_or_default().one_rtt as _)
    }),
    ("conn.rx.data-blocked", I, |_params, conn, _conns| {
        Some(conn.rx.unwrap_or_default().data_blocked as _)
    }),
    ("conn.rx.stream-data-blocked", I, |_params, conn, _conns| {
        Some(conn.rx.unwrap_or_default().stream_data_blocked as _)
    }),
    ("conn.rx.stream-progress", I, |_params, conn, _conns| {
        Some(conn.rx.unwrap_or_default().stream_progress as _)
    }),
    ("conn.rx.stream-throughput", Tpt, |_params, conn, _conns| {
        conn.rx?.stream_throughput()
    }),
    ("conn.lost.packets", I, |_params, conn, _conns| {
        Some(conn.loss.unwrap_or_default().packets() as _)
    }),
    ("conn.lost.initial", I, |_params, conn, _conns| {
        Some(conn.loss.unwrap_or_default().initial as _)
    }),
    ("conn.lost.handshake", I, |_params, conn, _conns| {
        Some(conn.loss.unwrap_or_default().handshake as _)
    }),
    ("conn.lost.one-rtt", I, |_params, conn, _conns| {
        Some(conn.loss.unwrap_or_default().one_rtt as _)
    }),
    ("sim.success", I, |_params, conn, conns| {
        // only return the value for the first connection
        if conn.id() == conns[0].id() {
            Some(conns.iter().filter(|c| c.is_success()).count() as _)
        } else {
            None
        }
    }),
    ("sim.error", I, |_params, conn, conns| {
        // only return the value for the first connection
        if conn.id() == conns[0].id() {
            Some(conns.iter().filter(|c| c.is_error()).count() as _)
        } else {
            None
        }
    }),
    ("net.drop-rate", P, |params, _conn, _conns| {
        Some(params.drop_rate)
    }),
    ("net.corrupt-rate", P, |params, _conn, _conns| {
        Some(params.corrupt_rate)
    }),
    ("net.jitter", T, |params, _conn, _conns| {
        Some(
            params
                .jitter
                .unwrap_or_default()
                .as_duration()
                .as_secs_f64(),
        )
    }),
    ("net.network-jitter", T, |params, _conn, _conns| {
        Some(
            params
                .network_jitter
                .unwrap_or_default()
                .as_duration()
                .as_secs_f64(),
        )
    }),
    ("net.delay", T, |params, _conn, _conns| {
        Some(params.delay.unwrap_or_default().as_duration().as_secs_f64())
    }),
    ("net.transmit-rate", I, |params, _conn, _conns| {
        Some(params.transmit_rate as f64)
    }),
    ("net.retransmit-rate", P, |params, _conn, _conns| {
        Some(params.retransmit_rate)
    }),
    ("net.max-udp-payload", I, |params, _conn, _conns| {
        Some(params.max_udp_payload as f64)
    }),
    ("net.max-inflight", I, |params, _conn, _conns| {
        Some(params.max_inflight as f64)
    }),
    ("net.endpoints", I, |params, _conn, _conns| {
        Some((params.servers + params.clients) as f64)
    }),
    ("net.servers", I, |params, _conn, _conns| {
        Some(params.servers as f64)
    }),
    ("net.clients", I, |params, _conn, _conns| {
        Some(params.clients as f64)
    }),
    ("net.connections", I, |_params, _conn, conns| {
        Some(conns.len() as f64)
    }),
    (
        "net.inflight-delay-threshold",
        I,
        |params, _conn, _conns| Some(params.inflight_delay_threshold as f64),
    ),
    ("net.inflight-delay", T, |params, _conn, _conns| {
        Some(
            params
                .inflight_delay
                .unwrap_or_default()
                .as_duration()
                .as_secs_f64(),
        )
    }),
];

pub static QUERY_NAMES: Lazy<Vec<&'static str>> =
    Lazy::new(|| QUERIES.iter().map(|q| q.0).collect());

#[derive(Clone, Copy)]
pub struct Query {
    pub name: &'static str,
    pub ty: Type,
    pub query: Q,
}

impl fmt::Debug for Query {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Query")
            .field("name", &self.name)
            .field("type", &self.ty)
            .finish()
    }
}

impl fmt::Display for Query {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.name.fmt(f)
    }
}

impl Query {
    pub fn apply(
        &self,
        params: &Parameters,
        conn: &Connection,
        conns: &[Connection],
    ) -> Option<f64> {
        (self.query)(params, conn, conns)
    }
}

impl FromStr for Query {
    type Err = anyhow::Error;

    fn from_str(path: &str) -> Result<Self, Self::Err> {
        for (name, ty, query) in QUERIES.iter().copied() {
            if name
                .split('-')
                .eq(path.split('_').flat_map(|v| v.split('-')))
            {
                return Ok(Self { name, ty, query });
            }
        }

        Err(anyhow::anyhow!("invalid query: {}", path))
    }
}

struct StrVisitor<T>(core::marker::PhantomData<T>);

impl<'de, T> serde::de::Visitor<'de> for StrVisitor<T>
where
    T: FromStr,
    <T as FromStr>::Err: core::fmt::Display,
{
    type Value = T;

    fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(formatter, "a valid thing")
    }

    fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
    where
        E: serde::de::Error,
    {
        v.parse().map_err(E::custom)
    }
}

impl<'de> serde::Deserialize<'de> for Query {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        deserializer.deserialize_str(StrVisitor::<Self>(Default::default()))
    }
}

type Op = fn(f64, f64) -> bool;

#[derive(Clone)]
pub struct Filter {
    pub expr: String,
    pub query: Query,
    pub value: f64,
    pub op: Op,
}

impl fmt::Debug for Filter {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Filter")
            .field("query", &self.query)
            .field("value", &self.value)
            .finish()
    }
}

impl fmt::Display for Filter {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.expr.fmt(f)
    }
}

impl Filter {
    pub fn apply(&self, params: &Parameters, conn: &Connection, conns: &[Connection]) -> bool {
        self.query
            .apply(params, conn, conns)
            .map_or(false, |actual| (self.op)(actual, self.value))
    }
}

impl FromStr for Filter {
    type Err = anyhow::Error;

    fn from_str(filter: &str) -> Result<Self, Self::Err> {
        let (query, op, value): (_, Op, _) = if let Some((q, v)) = filter.split_once("!=") {
            (q, |a, b| (a - b).abs() > f64::EPSILON, v)
        } else if let Some((q, v)) = filter.split_once("==") {
            (q, |a, b| (a - b).abs() < f64::EPSILON, v)
        } else if let Some((q, v)) = filter.split_once(">=") {
            (q, |a, b| a >= b, v)
        } else if let Some((q, v)) = filter.split_once('>') {
            (q, |a, b| a > b, v)
        } else if let Some((q, v)) = filter.split_once("<=") {
            (q, |a, b| a <= b, v)
        } else if let Some((q, v)) = filter.split_once('<') {
            (q, |a, b| a < b, v)
        } else if let Some((q, v)) = filter.split_once('=') {
            (q, |a, b| (a - b).abs() < f64::EPSILON, v)
        } else {
            (filter, |a, _b| a != 0.0, "0")
        };

        let query = Query::from_str(query)?;
        let value = query.ty.parse(value)?;

        Ok(Self {
            expr: filter.to_owned(),
            query,
            value,
            op,
        })
    }
}

impl<'de> serde::Deserialize<'de> for Filter {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        deserializer.deserialize_str(StrVisitor::<Self>(Default::default()))
    }
}


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

use super::CliRange;
use humantime::Duration;
use serde::Deserialize;
use structopt::StructOpt;

macro_rules! config {
    (struct Config { $(#[name = $name:literal] #[default = $default:literal] $field:ident: $ty:ty),* $(,)? }) => {
        mod defaults {
            use super::*;
            $(
            pub fn $field() -> $ty {
                $default.parse().unwrap()
            }
            )*
        }
        use defaults::*;

        #[derive(Clone, Debug, StructOpt, Deserialize)]
        #[serde(deny_unknown_fields)]
        pub struct Config {
            $(
                #[structopt(long, default_value = $default)]
                #[serde(default = $name)]
                pub $field: $ty,
            )*
        }

        impl Config {
            pub fn args(&self) -> Vec<String> {
                let mut args = vec![];

                $(
                    if self.$field != defaults::$field() {
                        args.push(format!("--{}", $name.replace('_', "-")));
                        args.push(self.$field.to_string());
                    }
                )*

                args
            }
        }
    };
}

config!(
    struct Config {
        #[name = "drop_rate"]
        #[default = "0.0"]
        drop_rate: CliRange<f64>,

        #[name = "corrupt_rate"]
        #[default = "0.0"]
        corrupt_rate: CliRange<f64>,

        #[name = "jitter"]
        #[default = "0ms"]
        jitter: CliRange<Duration>,

        #[name = "network_jitter"]
        #[default = "0ms"]
        network_jitter: CliRange<Duration>,

        #[name = "delay"]
        #[default = "100ms"]
        delay: CliRange<Duration>,

        #[name = "transmit_rate"]
        #[default = "0"]
        transmit_rate: CliRange<u64>,

        #[name = "retransmit_rate"]
        #[default = "0.0"]
        retransmit_rate: CliRange<f64>,

        #[name = "max_udp_payload"]
        #[default = "1450"]
        max_udp_payload: CliRange<u16>,

        #[name = "max_inflight"]
        #[default = "0"]
        max_inflight: CliRange<u64>,

        #[name = "inflight_delay"]
        #[default = "0ms"]
        inflight_delay: CliRange<Duration>,

        #[name = "inflight_delay_threshold"]
        #[default = "0"]
        inflight_delay_threshold: CliRange<u64>,

        #[name = "clients"]
        #[default = "1"]
        clients: CliRange<u32>,

        #[name = "servers"]
        #[default = "1"]
        servers: CliRange<u32>,

        #[name = "connect_delay"]
        #[default = "0ms"]
        connect_delay: CliRange<Duration>,

        #[name = "connections"]
        #[default = "1"]
        connections: CliRange<u32>,

        #[name = "streams"]
        #[default = "1"]
        streams: CliRange<u32>,

        #[name = "stream_data"]
        #[default = "4096"]
        stream_data: CliRange<u64>,

        #[name = "iterations"]
        #[default = "10000"]
        iterations: u64,
    }
);