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

use proc_macro2::TokenStream;
use quote::{quote, ToTokens};

type Error = Box<dyn std::error::Error>;
type Result<T, E = Error> = core::result::Result<T, E>;

mod parser;

#[derive(Debug, Default)]
struct Output {
    pub subscriber: TokenStream,
    pub endpoint_publisher: TokenStream,
    pub endpoint_publisher_subscriber: TokenStream,
    pub connection_publisher: TokenStream,
    pub connection_publisher_subscriber: TokenStream,
    pub tuple_subscriber: TokenStream,
    pub tracing_subscriber: TokenStream,
    pub builders: TokenStream,
    pub api: TokenStream,
    pub testing_fields: TokenStream,
    pub testing_fields_init: TokenStream,
    pub subscriber_testing: TokenStream,
    pub endpoint_publisher_testing: TokenStream,
    pub connection_publisher_testing: TokenStream,
    pub extra: TokenStream,
}

impl ToTokens for Output {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let Output {
            subscriber,
            endpoint_publisher,
            endpoint_publisher_subscriber,
            connection_publisher,
            connection_publisher_subscriber,
            tuple_subscriber,
            tracing_subscriber,
            builders,
            api,
            testing_fields,
            testing_fields_init,
            subscriber_testing,
            endpoint_publisher_testing,
            connection_publisher_testing,
            extra,
        } = self;

        tokens.extend(quote!(
            use super::*;

            pub mod api {
                //! This module contains events that are emitted to the [`Subscriber`](crate::event::Subscriber)
                use super::*;

                pub use traits::Subscriber;

                #api

                #extra
            }

            #[cfg(feature = "event-tracing")]
            pub mod tracing {
                //! This module contains event integration with [`tracing`](https://docs.rs/tracing)
                use super::api;

                /// Emits events with [`tracing`](https://docs.rs/tracing)
                #[derive(Clone, Debug)]
                pub struct Subscriber {
                    client: tracing::Span,
                    server: tracing::Span,
                }

                impl Default for Subscriber {
                    fn default() -> Self {
                        let root = tracing::span!(target: "s2n_quic", tracing::Level::DEBUG, "s2n_quic");
                        let client = tracing::span!(parent: root.id(), tracing::Level::DEBUG, "client");
                        let server = tracing::span!(parent: root.id(), tracing::Level::DEBUG, "server");

                        Self {
                            client,
                            server,
                        }
                    }
                }

                impl super::Subscriber for Subscriber {
                    type ConnectionContext = tracing::Span;

                    fn create_connection_context(&mut self, meta: &api::ConnectionMeta, _info: &api::ConnectionInfo) -> Self::ConnectionContext {
                        let parent = match meta.endpoint_type {
                            api::EndpointType::Client {} => {
                                self.client.id()
                            }
                            api::EndpointType::Server {} => {
                                self.server.id()
                            }
                        };
                        tracing::span!(target: "s2n_quic", parent: parent, tracing::Level::DEBUG, "conn", id = meta.id)
                    }

                    #tracing_subscriber
                }
            }

            pub mod builder {
                use super::*;

                #builders
            }

            pub mod supervisor {
                //! This module contains the `supervisor::Outcome` and `supervisor::Context` for use
                //! when implementing [`Subscriber::supervisor_timeout`](crate::event::Subscriber::supervisor_timeout) and
                //! [`Subscriber::on_supervisor_timeout`](crate::event::Subscriber::on_supervisor_timeout)
                //! on a Subscriber.

                use crate::{
                    application,
                    event::{builder::SocketAddress, IntoEvent},
                };

                #[non_exhaustive]
                #[derive(Clone, Debug, Eq, PartialEq)]
                pub enum Outcome {
                    /// Allow the connection to remain open
                    Continue,

                    /// Close the connection and notify the peer
                    Close {error_code: application::Error},

                    /// Close the connection without notifying the peer
                    ImmediateClose {reason: &'static str},
                }

                impl Default for Outcome {
                    fn default() -> Self {
                        Self::Continue
                    }
                }

                #[non_exhaustive]
                #[derive(Debug)]
                pub struct Context<'a> {
                    /// Number of handshakes that have begun but not completed
                    pub inflight_handshakes: usize,

                    /// Number of open connections
                    pub connection_count: usize,

                    /// The address of the peer
                    pub remote_address: SocketAddress<'a>,

                    /// True if the connection is in the handshake state, false otherwise
                    pub is_handshaking: bool,
                }

                impl<'a> Context<'a> {
                    pub fn new(
                        inflight_handshakes: usize,
                        connection_count: usize,
                        remote_address: &'a crate::inet::SocketAddress,
                        is_handshaking: bool,
                    ) -> Self {
                        Self {
                            inflight_handshakes,
                            connection_count,
                            remote_address: remote_address.into_event(),
                            is_handshaking,
                        }
                    }
                }
            }

            pub use traits::*;
            mod traits {
                use super::*;
                use api::*;
                use core::fmt;
                use crate::query;

                /// Provides metadata related to an event
                pub trait Meta: fmt::Debug {
                    /// Returns whether the local endpoint is a Client or Server
                    fn endpoint_type(&self) -> &EndpointType;

                    /// A context from which the event is being emitted
                    ///
                    /// An event can occur in the context of an Endpoint or Connection
                    fn subject(&self) -> Subject;

                    /// The time the event occurred
                    fn timestamp(&self) -> &crate::event::Timestamp;
                }

                impl Meta for ConnectionMeta {
                    fn endpoint_type(&self) -> &EndpointType {
                        &self.endpoint_type
                    }

                    fn subject(&self) -> Subject {
                        Subject::Connection { id : self.id }
                    }

                    fn timestamp(&self) -> &crate::event::Timestamp {
                        &self.timestamp
                    }
                }

                impl Meta for EndpointMeta {
                    fn endpoint_type(&self) -> &EndpointType {
                        &self.endpoint_type
                    }

                    fn subject(&self) -> Subject {
                        Subject::Endpoint {}
                    }

                    fn timestamp(&self) -> &crate::event::Timestamp {
                        &self.timestamp
                    }
                }

                /// Allows for events to be subscribed to
                pub trait Subscriber: 'static + Send {

                    /// An application provided type associated with each connection.
                    ///
                    /// The context provides a mechanism for applications to provide a custom type
                    /// and update it on each event, e.g. computing statistics. Each event
                    /// invocation (e.g. [`Subscriber::on_packet_sent`]) also provides mutable
                    /// access to the context `&mut ConnectionContext` and allows for updating the
                    /// context.
                    ///
                    /// ```no_run
                    /// # mod s2n_quic { pub mod provider { pub mod event {
                    /// #     pub use s2n_quic_core::event::{api as events, api::ConnectionInfo, api::ConnectionMeta, Subscriber};
                    /// # }}}
                    /// use s2n_quic::provider::event::{
                    ///     ConnectionInfo, ConnectionMeta, Subscriber, events::PacketSent
                    /// };
                    ///
                    /// pub struct MyEventSubscriber;
                    ///
                    /// pub struct MyEventContext {
                    ///     packet_sent: u64,
                    /// }
                    ///
                    /// impl Subscriber for MyEventSubscriber {
                    ///     type ConnectionContext = MyEventContext;
                    ///
                    ///     fn create_connection_context(
                    ///         &mut self, _meta: &ConnectionMeta,
                    ///         _info: &ConnectionInfo,
                    ///     ) -> Self::ConnectionContext {
                    ///         MyEventContext { packet_sent: 0 }
                    ///     }
                    ///
                    ///     fn on_packet_sent(
                    ///         &mut self,
                    ///         context: &mut Self::ConnectionContext,
                    ///         _meta: &ConnectionMeta,
                    ///         _event: &PacketSent,
                    ///     ) {
                    ///         context.packet_sent += 1;
                    ///     }
                    /// }
                    ///  ```
                    type ConnectionContext: 'static + Send;

                    /// Creates a context to be passed to each connection-related event
                    fn create_connection_context(&mut self, meta: &ConnectionMeta, info: &ConnectionInfo) -> Self::ConnectionContext;

                    /// The period at which `on_supervisor_timeout` is called
                    ///
                    /// If multiple `event::Subscriber`s are composed together, the minimum `supervisor_timeout`
                    /// across all `event::Subscriber`s will be used.
                    ///
                    /// If the `supervisor_timeout()` is `None` across all `event::Subscriber`s, connection supervision
                    /// will cease for the remaining lifetime of the connection and `on_supervisor_timeout` will no longer
                    /// be called.
                    ///
                    /// It is recommended to avoid setting this value less than ~100ms, as short durations
                    /// may lead to higher CPU utilization.
                    #[allow(unused_variables)]
                    fn supervisor_timeout(&mut self, conn_context: &mut Self::ConnectionContext, meta: &ConnectionMeta, context: &supervisor::Context) -> Option<Duration> {
                        None
                    }

                    /// Called for each `supervisor_timeout` to determine any action to take on the connection based on the `supervisor::Outcome`
                    ///
                    /// If multiple `event::Subscriber`s are composed together, the minimum `supervisor_timeout`
                    /// across all `event::Subscriber`s will be used, and thus `on_supervisor_timeout` may be called
                    /// earlier than the `supervisor_timeout` for a given `event::Subscriber` implementation.
                    #[allow(unused_variables)]
                    fn on_supervisor_timeout(&mut self, conn_context: &mut Self::ConnectionContext, meta: &ConnectionMeta, context: &supervisor::Context) -> supervisor::Outcome {
                        supervisor::Outcome::default()
                    }

                    #subscriber

                    /// Called for each event that relates to the endpoint and all connections
                    #[inline]
                    fn on_event<M: Meta, E: Event>(&mut self, meta: &M, event: &E) {
                        let _ = meta;
                        let _ = event;
                    }

                    /// Called for each event that relates to a connection
                    #[inline]
                    fn on_connection_event<E: Event>(&mut self, context: &mut Self::ConnectionContext, meta: &ConnectionMeta, event: &E) {
                        let _ = context;
                        let _ = meta;
                        let _ = event;
                    }

                    /// Used for querying the `Subscriber::ConnectionContext` on a Subscriber
                    #[inline]
                    fn query(context: &Self::ConnectionContext, query: &mut dyn query::Query) -> query::ControlFlow {
                        query.execute(context)
                    }

                    /// Used for querying and mutating the `Subscriber::ConnectionContext` on a Subscriber
                    #[inline]
                    fn query_mut(context: &mut Self::ConnectionContext, query: &mut dyn query::QueryMut) -> query::ControlFlow {
                        query.execute_mut(context)
                    }
                }

                /// Subscriber is implemented for a 2-element tuple to make it easy to compose multiple
                /// subscribers.
                impl<A, B> Subscriber for (A, B)
                    where
                        A: Subscriber,
                        B: Subscriber,
                {
                    type ConnectionContext = (A::ConnectionContext, B::ConnectionContext);

                    #[inline]
                    fn create_connection_context(&mut self, meta: &ConnectionMeta, info: &ConnectionInfo) -> Self::ConnectionContext {
                        (self.0.create_connection_context(meta, info), self.1.create_connection_context(meta, info))
                    }

                    #[inline]
                    fn supervisor_timeout(&mut self, conn_context: &mut Self::ConnectionContext, meta: &ConnectionMeta, context: &supervisor::Context) -> Option<Duration> {
                        let timeout_a = self.0.supervisor_timeout(&mut conn_context.0, meta, context);
                        let timeout_b = self.1.supervisor_timeout(&mut conn_context.1, meta, context);
                        match (timeout_a, timeout_b) {
                            (None, None) => None,
                            (None, Some(timeout)) | (Some(timeout), None) => Some(timeout),
                            (Some(a), Some(b)) => Some(a.min(b)),
                        }
                    }

                    #[inline]
                    fn on_supervisor_timeout(&mut self, conn_context: &mut Self::ConnectionContext, meta: &ConnectionMeta, context: &supervisor::Context) -> supervisor::Outcome {
                        let outcome_a = self.0.on_supervisor_timeout(&mut conn_context.0, meta, context);
                        let outcome_b = self.1.on_supervisor_timeout(&mut conn_context.1, meta, context);
                        match (outcome_a, outcome_b) {
                            (supervisor::Outcome::ImmediateClose { reason }, _) | (_, supervisor::Outcome::ImmediateClose { reason }) => supervisor::Outcome::ImmediateClose { reason },
                            (supervisor::Outcome::Close { error_code }, _) | (_, supervisor::Outcome::Close { error_code }) => supervisor::Outcome::Close { error_code },
                            _ => supervisor::Outcome::Continue,
                        }
                    }

                    #tuple_subscriber

                    #[inline]
                    fn on_event<M: Meta, E: Event>(&mut self, meta: &M, event: &E) {
                        self.0.on_event(meta, event);
                        self.1.on_event(meta, event);
                    }

                    #[inline]
                    fn on_connection_event<E: Event>(&mut self, context: &mut Self::ConnectionContext, meta: &ConnectionMeta, event: &E) {
                        self.0.on_connection_event(&mut context.0, meta, event);
                        self.1.on_connection_event(&mut context.1, meta, event);
                    }

                    #[inline]
                    fn query(context: &Self::ConnectionContext, query: &mut dyn query::Query) -> query::ControlFlow {
                        query.execute(context)
                            .and_then(|| A::query(&context.0, query))
                            .and_then(|| B::query(&context.1, query))
                    }

                    #[inline]
                    fn query_mut(context: &mut Self::ConnectionContext, query: &mut dyn query::QueryMut) -> query::ControlFlow {
                        query.execute_mut(context)
                            .and_then(|| A::query_mut(&mut context.0, query))
                            .and_then(|| B::query_mut(&mut context.1, query))
                    }
                }

                pub trait EndpointPublisher {
                    #endpoint_publisher

                    /// Returns the QUIC version, if any
                    fn quic_version(&self) -> Option<u32>;
                }

                pub struct EndpointPublisherSubscriber<'a, Sub: Subscriber> {
                    meta: EndpointMeta,
                    quic_version: Option<u32>,
                    subscriber: &'a mut Sub,
                }

                impl<'a, Sub: Subscriber> fmt::Debug for EndpointPublisherSubscriber<'a, Sub> {
                    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
                        f.debug_struct("ConnectionPublisherSubscriber")
                            .field("meta", &self.meta)
                            .field("quic_version", &self.quic_version)
                            .finish()
                    }
                }

                impl<'a, Sub: Subscriber> EndpointPublisherSubscriber<'a, Sub> {
                    #[inline]
                    pub fn new(
                        meta: builder::EndpointMeta,
                        quic_version: Option<u32>,
                        subscriber: &'a mut Sub,
                    ) -> Self {
                        Self {
                            meta: meta.into_event(),
                            quic_version,
                            subscriber,
                        }
                    }
                }

                impl<'a, Sub: Subscriber> EndpointPublisher for EndpointPublisherSubscriber<'a, Sub> {
                    #endpoint_publisher_subscriber

                    #[inline]
                    fn quic_version(&self) -> Option<u32> {
                        self.quic_version
                    }
                }

                pub trait ConnectionPublisher {
                    #connection_publisher

                    /// Returns the QUIC version negotiated for the current connection, if any
                    fn quic_version(&self) -> u32;

                    /// Returns the [`Subject`] for the current publisher
                    fn subject(&self) -> Subject;
                }

                pub struct ConnectionPublisherSubscriber<'a, Sub: Subscriber> {
                    meta: ConnectionMeta,
                    quic_version: u32,
                    subscriber: &'a mut Sub,
                    context: &'a mut Sub::ConnectionContext,
                }

                impl<'a, Sub: Subscriber> fmt::Debug for ConnectionPublisherSubscriber<'a, Sub> {
                    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
                        f.debug_struct("ConnectionPublisherSubscriber")
                            .field("meta", &self.meta)
                            .field("quic_version", &self.quic_version)
                            .finish()
                    }
                }

                impl<'a, Sub: Subscriber> ConnectionPublisherSubscriber<'a, Sub> {
                    #[inline]
                    pub fn new(
                        meta: builder::ConnectionMeta,
                        quic_version: u32,
                        subscriber: &'a mut Sub,
                        context: &'a mut Sub::ConnectionContext
                    ) -> Self {
                        Self {
                            meta: meta.into_event(),
                            quic_version,
                            subscriber,
                            context,
                        }
                    }
                }

                impl<'a, Sub: Subscriber> ConnectionPublisher for ConnectionPublisherSubscriber<'a, Sub> {
                    #connection_publisher_subscriber

                    #[inline]
                    fn quic_version(&self) -> u32 {
                        self.quic_version
                    }

                    #[inline]
                    fn subject(&self) -> api::Subject {
                        self.meta.subject()
                    }
                }
            }

            #[cfg(any(test, feature = "testing"))]
            pub mod testing {
                use super::*;

                #[derive(Clone, Debug)]
                pub struct Subscriber {
                    location: Option<Location>,
                    output: Vec<String>,
                    #testing_fields
                }

                impl Drop for Subscriber {
                    fn drop(&mut self) {
                        // don't make any assertions if we're already failing the test
                        if std::thread::panicking() {
                            return;
                        }

                        if let Some(location) = self.location.as_ref() {
                            location.snapshot(&self.output);
                        }
                    }
                }

                impl Subscriber {
                    /// Creates a subscriber with snapshot assertions enabled
                    #[track_caller]
                    pub fn snapshot() -> Self {
                        let mut sub = Self::no_snapshot();
                        sub.location = Location::try_new();
                        sub
                    }

                    /// Creates a subscriber with snapshot assertions disabled
                    pub fn no_snapshot() -> Self {
                        Self {
                            location: None,
                            output: Default::default(),
                            #testing_fields_init
                        }
                    }
                }

                impl super::Subscriber for Subscriber {
                    type ConnectionContext = ();

                    fn create_connection_context(&mut self, _meta: &api::ConnectionMeta, _info: &api::ConnectionInfo) -> Self::ConnectionContext {}

                    #subscriber_testing
                }

                #[derive(Clone, Debug)]
                pub struct Publisher {
                    location: Option<Location>,
                    output: Vec<String>,
                    #testing_fields
                }

                impl Publisher {
                    /// Creates a publisher with snapshot assertions enabled
                    #[track_caller]
                    pub fn snapshot() -> Self {
                        let mut sub = Self::no_snapshot();
                        sub.location = Location::try_new();
                        sub
                    }

                    /// Creates a publisher with snapshot assertions disabled
                    pub fn no_snapshot() -> Self {
                        Self {
                            location: None,
                            output: Default::default(),
                            #testing_fields_init
                        }
                    }
                }

                impl super::EndpointPublisher for Publisher {
                    #endpoint_publisher_testing

                    fn quic_version(&self) -> Option<u32> {
                        Some(1)
                    }
                }

                impl super::ConnectionPublisher for Publisher {
                    #connection_publisher_testing

                    fn quic_version(&self) -> u32 {
                        1
                    }

                    fn subject(&self) -> api::Subject {
                        api::Subject::Connection { id: 0 }
                    }
                }

                impl Drop for Publisher {
                    fn drop(&mut self) {
                        // don't make any assertions if we're already failing the test
                        if std::thread::panicking() {
                            return;
                        }

                        if let Some(location) = self.location.as_ref() {
                            location.snapshot(&self.output);
                        }
                    }
                }

                #[derive(Clone, Debug)]
                struct Location(&'static core::panic::Location<'static>);

                impl Location {
                    #[track_caller]
                    fn try_new() -> Option<Self> {
                        let thread = std::thread::current();

                        // only create a location if insta can figure out the test name from the
                        // thread
                        if thread.name().map_or(false, |name| name != "main") {
                            Some(Self(core::panic::Location::caller()))
                        } else {
                            None
                        }
                    }

                    fn snapshot(&self, output: &[String]) {
                        // miri doesn't support the syscalls that insta uses
                        if cfg!(miri) {
                            return;
                        }

                        use std::path::{Path, Component};
                        let value = output.join("\n");

                        let thread = std::thread::current();
                        let function_name = thread.name().unwrap();

                        let test_path = Path::new(self.0.file().trim_end_matches(".rs"));
                        let module_path = test_path
                            .components()
                            .filter_map(|comp| match comp {
                                Component::Normal(comp) => comp.to_str(),
                                _ => Some("_"),
                            })
                            .chain(Some("events"))
                            .collect::<Vec<_>>()
                            .join("::");

                        let current_dir = std::env::current_dir().unwrap();

                        insta::_macro_support::assert_snapshot(
                            insta::_macro_support::AutoName.into(),
                            &value,
                            current_dir.to_str().unwrap(),
                            function_name,
                            &module_path,
                            self.0.file(),
                            self.0.line(),
                            "",
                        )
                        .unwrap()
                    }
                }
            }
        ));
    }
}

fn main() -> Result<()> {
    let mut files = vec![];

    for path in glob::glob(concat!(env!("CARGO_MANIFEST_DIR"), "/events/**/*.rs"))? {
        let path = path?;
        let file = std::fs::read_to_string(path)?;
        files.push(parser::parse(&file).unwrap());
    }

    let mut output = Output::default();

    for file in &files {
        file.to_tokens(&mut output);
    }

    let generated = concat!(
        env!("CARGO_MANIFEST_DIR"),
        "/../s2n-quic-core/src/event/generated.rs"
    );

    let mut o = std::fs::File::create(generated)?;

    macro_rules! put {
        ($($arg:tt)*) => {{
            use std::io::Write;
            writeln!(o, $($arg)*)?;
        }}
    }

    put!("// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.");
    put!("// SPDX-License-Identifier: Apache-2.0");
    put!();
    put!("// DO NOT MODIFY THIS FILE");
    put!("// This file was generated with the `s2n-quic-events` crate and any required");
    put!("// changes should be made there.");
    put!();
    put!("{}", output.to_token_stream());

    let status = std::process::Command::new("rustfmt")
        .arg(generated)
        .spawn()?
        .wait()?;

    assert!(status.success());

    Ok(())
}


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

use crate::{Output, Result};
use heck::ToSnakeCase;
use proc_macro2::{Ident, Span, TokenStream};
use quote::{quote, ToTokens};
use syn::{
    parse::{Parse, ParseStream},
    Meta,
};

#[derive(Debug, Default)]
pub struct File {
    pub structs: Vec<Struct>,
    pub enums: Vec<Enum>,
    pub extra: TokenStream,
}

impl File {
    fn parse(file: syn::File) -> Self {
        assert!(file.attrs.is_empty());
        assert!(file.shebang.is_none());

        let mut out = File::default();
        for item in file.items {
            match item {
                syn::Item::Enum(v) => out.enums.push(Enum::parse(v)),
                syn::Item::Struct(v) => out.structs.push(Struct::parse(v)),
                item => item.to_tokens(&mut out.extra),
            }
        }

        out
    }

    pub(crate) fn to_tokens(&self, output: &mut Output) {
        self.extra.to_tokens(&mut output.extra);
        for v in &self.structs {
            v.to_tokens(output);
        }
        for v in &self.enums {
            v.to_tokens(output);
        }
    }
}

#[derive(Debug)]
pub struct Struct {
    attrs: ContainerAttrs,
    ident: syn::Ident,
    generics: syn::Generics,
    fields: Vec<Field>,
}

impl Struct {
    fn parse(item: syn::ItemStruct) -> Self {
        let attrs = ContainerAttrs::parse(item.attrs);

        Self {
            attrs,
            ident: item.ident,
            generics: item.generics,
            fields: item.fields.into_iter().map(Field::parse).collect(),
        }
    }

    fn to_tokens(&self, output: &mut Output) {
        let Self {
            attrs,
            ident,
            generics,
            fields,
        } = self;

        let derive_attrs = &attrs.derive_attrs;
        let builder_derive_attrs = &attrs.builder_derive_attrs;
        let extra_attrs = &attrs.extra;
        let deprecated = &attrs.deprecated;
        let allow_deprecated = &attrs.allow_deprecated;

        let destructure_fields: Vec<_> = fields.iter().map(Field::destructure).collect();
        let builder_fields = fields.iter().map(Field::builder);
        let builder_field_impls = fields.iter().map(Field::builder_impl);
        let api_fields = fields.iter().map(Field::api);

        if attrs.builder_derive {
            output.builders.extend(quote!(
                #[#builder_derive_attrs]
            ));
        }

        output.builders.extend(quote!(
            #[derive(Clone, Debug)]
            #extra_attrs
            pub struct #ident #generics {
                #(#builder_fields)*
            }

            #allow_deprecated
            impl #generics IntoEvent<api::#ident #generics> for #ident #generics {
                #[inline]
                fn into_event(self) -> api::#ident #generics {
                    let #ident {
                        #(#destructure_fields),*
                    } = self;

                    api::#ident {
                        #(#builder_field_impls)*
                    }
                }
            }
        ));

        if attrs.derive {
            output.api.extend(quote!(#[derive(Clone, Debug)]));
        }

        if !attrs.exhaustive {
            output.api.extend(quote!(#[non_exhaustive]));
        }

        output.api.extend(quote!(
            #derive_attrs
            #extra_attrs
            #deprecated
            #allow_deprecated
            pub struct #ident #generics {
                #(#api_fields)*
            }
        ));

        if let Some(event_name) = attrs.event_name.as_ref() {
            output.api.extend(quote!(
                #allow_deprecated
                impl #generics Event for #ident #generics {
                    const NAME: &'static str = #event_name;
                }
            ));

            let ident_str = ident.to_string();
            let snake = ident_str.to_snake_case();
            let function = format!("on_{snake}");
            let counter = Ident::new(&snake, Span::call_site());
            let function = Ident::new(&function, Span::call_site());

            let subscriber_doc = format!("Called when the `{ident_str}` event is triggered");
            let publisher_doc =
                format!("Publishes a `{ident_str}` event to the publisher's subscriber");

            // add a counter for testing structs
            output.testing_fields.extend(quote!(
                pub #counter: u32,
            ));
            output.testing_fields_init.extend(quote!(
                #counter: 0,
            ));

            match attrs.subject {
                Subject::Endpoint => {
                    output.subscriber.extend(quote!(
                        #[doc = #subscriber_doc]
                        #[inline]
                        #deprecated
                        #allow_deprecated
                        fn #function(&mut self, meta: &EndpointMeta, event: &#ident) {
                            let _ = meta;
                            let _ = event;
                        }
                    ));

                    output.tuple_subscriber.extend(quote!(
                        #[inline]
                        #allow_deprecated
                        fn #function(&mut self, meta: &EndpointMeta, event: &#ident) {
                            (self.0).#function(meta, event);
                            (self.1).#function(meta, event);
                        }
                    ));

                    output.tracing_subscriber.extend(quote!(
                        #[inline]
                        #allow_deprecated
                        fn #function(&mut self, meta: &api::EndpointMeta, event: &api::#ident) {
                            let parent = match meta.endpoint_type {
                                api::EndpointType::Client {} => {
                                    self.client.id()
                                }
                                api::EndpointType::Server {} => {
                                    self.server.id()
                                }
                            };
                            let api::#ident { #(#destructure_fields),* } = event;
                            tracing::event!(target: #snake, parent: parent, tracing::Level::DEBUG, #(#destructure_fields = tracing::field::debug(#destructure_fields)),*);
                        }
                    ));

                    output.endpoint_publisher.extend(quote!(
                        #[doc = #publisher_doc]
                        fn #function(&mut self, event: builder::#ident);
                    ));

                    output.endpoint_publisher_subscriber.extend(quote!(
                        #[inline]
                        #allow_deprecated
                        fn #function(&mut self, event: builder::#ident) {
                            let event = event.into_event();
                            self.subscriber.#function(&self.meta, &event);
                            self.subscriber.on_event(&self.meta, &event);
                        }
                    ));

                    output.subscriber_testing.extend(quote!(
                        #allow_deprecated
                        fn #function(&mut self, meta: &api::EndpointMeta, event: &api::#ident) {
                            self.#counter += 1;
                            self.output.push(format!("{meta:?} {event:?}"));
                        }
                    ));

                    output.endpoint_publisher_testing.extend(quote!(
                        #allow_deprecated
                        fn #function(&mut self, event: builder::#ident) {
                            self.#counter += 1;
                            let event = event.into_event();
                            self.output.push(format!("{event:?}"));
                        }
                    ));
                }
                Subject::Connection => {
                    output.subscriber.extend(quote!(
                        #[doc = #subscriber_doc]
                        #[inline]
                        #deprecated
                        #allow_deprecated
                        fn #function(&mut self, context: &mut Self::ConnectionContext, meta: &ConnectionMeta, event: &#ident) {
                            let _ = context;
                            let _ = meta;
                            let _ = event;
                        }
                    ));

                    output.tuple_subscriber.extend(quote!(
                        #[inline]
                        #allow_deprecated
                        fn #function(&mut self, context: &mut Self::ConnectionContext, meta: &ConnectionMeta, event: &#ident) {
                            (self.0).#function(&mut context.0, meta, event);
                            (self.1).#function(&mut context.1, meta, event);
                        }
                    ));

                    output.tracing_subscriber.extend(quote!(
                        #[inline]
                        #allow_deprecated
                        fn #function(&mut self, context: &mut Self::ConnectionContext, _meta: &api::ConnectionMeta, event: &api::#ident) {
                            let id = context.id();
                            let api::#ident { #(#destructure_fields),* } = event;
                            tracing::event!(target: #snake, parent: id, tracing::Level::DEBUG, #(#destructure_fields = tracing::field::debug(#destructure_fields)),*);
                        }
                    ));

                    output.connection_publisher.extend(quote!(
                        #[doc = #publisher_doc]
                        fn #function(&mut self, event: builder::#ident);
                    ));

                    output.connection_publisher_subscriber.extend(quote!(
                        #[inline]
                        #allow_deprecated
                        fn #function(&mut self, event: builder::#ident) {
                            let event = event.into_event();
                            self.subscriber.#function(self.context, &self.meta, &event);
                            self.subscriber.on_connection_event(self.context, &self.meta, &event);
                            self.subscriber.on_event(&self.meta, &event);
                        }
                    ));

                    output.subscriber_testing.extend(quote!(
                        #allow_deprecated
                        fn #function(&mut self, _context: &mut Self::ConnectionContext, meta: &api::ConnectionMeta, event: &api::#ident) {
                            self.#counter += 1;
                            if self.location.is_some() {
                                self.output.push(format!("{meta:?} {event:?}"));
                            }
                        }
                    ));

                    output.connection_publisher_testing.extend(quote!(
                        #allow_deprecated
                        fn #function(&mut self, event: builder::#ident) {
                            self.#counter += 1;
                            let event = event.into_event();
                            if self.location.is_some() {
                                self.output.push(format!("{event:?}"));
                            }
                        }
                    ));
                }
            }
        }
    }
}

#[derive(Debug)]
pub struct Enum {
    attrs: ContainerAttrs,
    ident: syn::Ident,
    generics: syn::Generics,
    variants: Vec<Variant>,
}

impl Enum {
    fn parse(item: syn::ItemEnum) -> Self {
        let attrs = ContainerAttrs::parse(item.attrs);

        Self {
            attrs,
            ident: item.ident,
            generics: item.generics,
            variants: item.variants.into_iter().map(Variant::parse).collect(),
        }
    }

    fn to_tokens(&self, output: &mut Output) {
        let Self {
            attrs,
            ident,
            generics,
            variants,
        } = self;

        assert!(
            attrs.event_name.is_none(),
            "enum events are not currently supported"
        );

        let derive_attrs = &attrs.derive_attrs;
        let builder_derive_attrs = &attrs.builder_derive_attrs;
        let extra_attrs = &attrs.extra;
        let deprecated = &attrs.deprecated;
        let allow_deprecated = &attrs.allow_deprecated;

        let builder_fields = variants.iter().map(Variant::builder);
        let builder_field_impls = variants.iter().map(Variant::builder_impl);
        let api_fields = variants.iter().map(Variant::api);

        if attrs.builder_derive {
            output.builders.extend(quote!(
                #[#builder_derive_attrs]
            ));
        }

        output.builders.extend(quote!(
            #[derive(Clone, Debug)]
            #extra_attrs
            pub enum #ident #generics {
                #(#builder_fields)*
            }

            #allow_deprecated
            impl #generics IntoEvent<api::#ident #generics> for #ident #generics {
                #[inline]
                fn into_event(self) -> api::#ident #generics {
                    use api::#ident::*;
                    match self {
                        #(#builder_field_impls)*
                    }
                }
            }
        ));

        if attrs.derive {
            output.api.extend(quote!(#[derive(Clone, Debug)]));
        }

        if !attrs.exhaustive {
            output.api.extend(quote!(#[non_exhaustive]));
        }

        output.api.extend(quote!(
            #derive_attrs
            #extra_attrs
            #deprecated
            pub enum #ident #generics {
                #(#api_fields)*
            }
        ));
    }
}

#[derive(Debug)]
struct ContainerAttrs {
    event_name: Option<syn::LitStr>,
    deprecated: TokenStream,
    allow_deprecated: TokenStream,
    subject: Subject,
    exhaustive: bool,
    derive: bool,
    derive_attrs: TokenStream,
    builder_derive: bool,
    builder_derive_attrs: TokenStream,
    extra: TokenStream,
}

impl ContainerAttrs {
    fn parse(attrs: Vec<syn::Attribute>) -> Self {
        let mut v = Self {
            // events must include a name to be considered an event
            event_name: None,
            deprecated: TokenStream::default(),
            allow_deprecated: TokenStream::default(),
            // most event subjects relate to actual connections so make that the default
            subject: Subject::Connection,
            // default to #[non_exhaustive]
            exhaustive: false,
            derive: true,
            derive_attrs: quote!(),
            builder_derive: false,
            builder_derive_attrs: quote!(),
            extra: quote!(),
        };

        for attr in attrs {
            let path = attr.path();
            if path.is_ident("event") {
                v.event_name = Some(attr.parse_args().unwrap());
            } else if path.is_ident("deprecated") {
                attr.to_tokens(&mut v.deprecated);

                if v.allow_deprecated.is_empty() {
                    v.allow_deprecated = quote!(#[allow(deprecated)]);
                }
            } else if path.is_ident("subject") {
                v.subject = attr.parse_args().unwrap();
            } else if path.is_ident("exhaustive") {
                v.exhaustive = true;
            } else if path.is_ident("derive") {
                v.derive = false;
                attr.to_tokens(&mut v.derive_attrs);
            } else if path.is_ident("builder_derive") {
                v.builder_derive = true;
                if let Meta::List(list) = attr.parse_args().unwrap() {
                    list.to_tokens(&mut v.builder_derive_attrs);
                }
            } else {
                attr.to_tokens(&mut v.extra)
            }
        }

        v
    }
}

#[derive(Debug)]
enum Subject {
    Connection,
    Endpoint,
}

impl Parse for Subject {
    fn parse(input: ParseStream) -> syn::parse::Result<Self> {
        let name: Ident = input.parse()?;
        match name.to_string().as_str() {
            "connection" => Ok(Self::Connection),
            "endpoint" => Ok(Self::Endpoint),
            name => Err(syn::parse::Error::new(
                input.span(),
                format!("invalid event subject: {name}, expected connection or endpoint"),
            )),
        }
    }
}

#[derive(Debug)]
struct Field {
    attrs: FieldAttrs,
    ident: Option<syn::Ident>,
    ty: syn::Type,
}

impl Field {
    fn parse(item: syn::Field) -> Self {
        Self {
            attrs: FieldAttrs::parse(item.attrs),
            ident: item.ident,
            ty: item.ty,
        }
    }

    fn api(&self) -> TokenStream {
        let Self { attrs, ident, ty } = self;
        let attrs = &attrs.extra;
        if let Some(name) = ident {
            quote!(
                #attrs
                pub #name: #ty,
            )
        } else {
            quote!(#attrs pub #ty,)
        }
    }

    fn enum_api(&self) -> TokenStream {
        let Self {
            attrs, ident, ty, ..
        } = self;
        let attrs = &attrs.extra;
        if let Some(name) = ident {
            quote!(
                #attrs
                #name: #ty,
            )
        } else {
            quote!(#attrs #ty,)
        }
    }

    fn destructure(&self) -> TokenStream {
        let Self { ident, .. } = self;
        quote!(#ident)
    }

    fn builder(&self) -> TokenStream {
        let Self { attrs, ident, .. } = self;
        let attrs = &attrs.extra;
        let ty = self.builder_type();
        if let Some(name) = ident {
            quote!(#attrs pub #name: #ty,)
        } else {
            quote!(#attrs pub #ty,)
        }
    }

    fn enum_builder(&self) -> TokenStream {
        let Self { attrs, ident, .. } = self;
        let attrs = &attrs.extra;
        let ty = self.builder_type();
        if let Some(name) = ident {
            quote!(#attrs #name: #ty,)
        } else {
            quote!(#attrs #ty,)
        }
    }

    fn builder_impl(&self) -> TokenStream {
        let Self { ident, .. } = self;
        quote!(#ident: #ident.into_event(),)
    }

    fn builder_type(&self) -> &syn::Type {
        if let Some(ty) = &self.attrs.builder {
            ty
        } else {
            &self.ty
        }
    }
}

#[derive(Debug)]
struct FieldAttrs {
    builder: Option<syn::Type>,
    extra: TokenStream,
}

impl FieldAttrs {
    fn parse(attrs: Vec<syn::Attribute>) -> Self {
        let mut v = Self {
            // The event can override the builder with a specific type
            builder: None,
            extra: quote!(),
        };

        for attr in attrs {
            if attr.path().is_ident("builder") {
                v.builder = Some(attr.parse_args().unwrap());
            } else {
                attr.to_tokens(&mut v.extra)
            }
        }

        v
    }
}

#[derive(Debug)]
struct Variant {
    ident: syn::Ident,
    attrs: Vec<syn::Attribute>,
    fields: Vec<Field>,
}

impl Variant {
    fn parse(item: syn::Variant) -> Self {
        Self {
            ident: item.ident,
            attrs: item.attrs,
            fields: item.fields.into_iter().map(Field::parse).collect(),
        }
    }

    fn api(&self) -> TokenStream {
        let Self {
            ident,
            attrs,
            fields,
        } = self;
        let fields = fields.iter().map(Field::enum_api);
        quote!(
            #[non_exhaustive]
            #(#attrs)*
            #ident { #(#fields)* },
        )
    }

    fn builder(&self) -> TokenStream {
        let Self {
            ident,
            fields,
            attrs,
        } = self;
        if fields.is_empty() {
            quote!(#(#attrs)* #ident,)
        } else {
            let fields = fields.iter().map(Field::enum_builder);
            quote!(#(#attrs)* #ident { #(#fields)* },)
        }
    }

    fn builder_impl(&self) -> TokenStream {
        let Self { ident, fields, .. } = self;
        if fields.is_empty() {
            quote!(Self::#ident => #ident { },)
        } else {
            let destructure = fields.iter().map(Field::destructure);
            let fields = fields.iter().map(Field::builder_impl);

            quote!(Self::#ident { #(#destructure),* } => #ident { #(#fields)* },)
        }
    }
}

pub fn parse(contents: &str) -> Result<File> {
    let file = syn::parse_str(contents)?;
    let common = File::parse(file);
    Ok(common)
}


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

mod s2n_quic;

pub use self::s2n_quic::*;
pub use h3;


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

use bytes::{Buf, Bytes};
use futures::ready;
use h3::quic::{self, Error, StreamId, WriteBuf};
use s2n_quic::stream::{BidirectionalStream, ReceiveStream};
use s2n_quic_core::varint::VarInt;
use std::{
    convert::TryInto,
    fmt::{self, Display},
    sync::Arc,
    task::{self, Poll},
};

pub struct Connection {
    conn: s2n_quic::connection::Handle,
    bidi_acceptor: s2n_quic::connection::BidirectionalStreamAcceptor,
    recv_acceptor: s2n_quic::connection::ReceiveStreamAcceptor,
}

impl Connection {
    pub fn new(new_conn: s2n_quic::Connection) -> Self {
        let (handle, acceptor) = new_conn.split();
        let (bidi, recv) = acceptor.split();

        Self {
            conn: handle,
            bidi_acceptor: bidi,
            recv_acceptor: recv,
        }
    }
}

#[derive(Debug)]
pub struct ConnectionError(s2n_quic::connection::Error);

impl std::error::Error for ConnectionError {}

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

impl Error for ConnectionError {
    fn is_timeout(&self) -> bool {
        matches!(self.0, s2n_quic::connection::Error::IdleTimerExpired { .. })
    }

    fn err_code(&self) -> Option<u64> {
        match self.0 {
            s2n_quic::connection::Error::Application { error, .. } => Some(error.into()),
            _ => None,
        }
    }
}

impl From<s2n_quic::connection::Error> for ConnectionError {
    fn from(e: s2n_quic::connection::Error) -> Self {
        Self(e)
    }
}

impl<B> quic::Connection<B> for Connection
where
    B: Buf,
{
    type BidiStream = BidiStream<B>;
    type SendStream = SendStream<B>;
    type RecvStream = RecvStream;
    type OpenStreams = OpenStreams;
    type Error = ConnectionError;

    fn poll_accept_recv(
        &mut self,
        cx: &mut task::Context<'_>,
    ) -> Poll<Result<Option<Self::RecvStream>, Self::Error>> {
        let recv = match ready!(self.recv_acceptor.poll_accept_receive_stream(cx))? {
            Some(x) => x,
            None => return Poll::Ready(Ok(None)),
        };
        Poll::Ready(Ok(Some(Self::RecvStream::new(recv))))
    }

    fn poll_accept_bidi(
        &mut self,
        cx: &mut task::Context<'_>,
    ) -> Poll<Result<Option<Self::BidiStream>, Self::Error>> {
        let (recv, send) = match ready!(self.bidi_acceptor.poll_accept_bidirectional_stream(cx))? {
            Some(x) => x.split(),
            None => return Poll::Ready(Ok(None)),
        };
        Poll::Ready(Ok(Some(Self::BidiStream {
            send: Self::SendStream::new(send),
            recv: Self::RecvStream::new(recv),
        })))
    }

    fn poll_open_bidi(
        &mut self,
        cx: &mut task::Context<'_>,
    ) -> Poll<Result<Self::BidiStream, Self::Error>> {
        let stream = ready!(self.conn.poll_open_bidirectional_stream(cx))?;
        Ok(stream.into()).into()
    }

    fn poll_open_send(
        &mut self,
        cx: &mut task::Context<'_>,
    ) -> Poll<Result<Self::SendStream, Self::Error>> {
        let stream = ready!(self.conn.poll_open_send_stream(cx))?;
        Ok(stream.into()).into()
    }

    fn opener(&self) -> Self::OpenStreams {
        OpenStreams {
            conn: self.conn.clone(),
        }
    }

    fn close(&mut self, code: h3::error::Code, _reason: &[u8]) {
        self.conn.close(
            code.value()
                .try_into()
                .expect("s2n-quic supports error codes up to 2^62-1"),
        );
    }
}

pub struct OpenStreams {
    conn: s2n_quic::connection::Handle,
}

impl<B> quic::OpenStreams<B> for OpenStreams
where
    B: Buf,
{
    type BidiStream = BidiStream<B>;
    type SendStream = SendStream<B>;
    type RecvStream = RecvStream;
    type Error = ConnectionError;

    fn poll_open_bidi(
        &mut self,
        cx: &mut task::Context<'_>,
    ) -> Poll<Result<Self::BidiStream, Self::Error>> {
        let stream = ready!(self.conn.poll_open_bidirectional_stream(cx))?;
        Ok(stream.into()).into()
    }

    fn poll_open_send(
        &mut self,
        cx: &mut task::Context<'_>,
    ) -> Poll<Result<Self::SendStream, Self::Error>> {
        let stream = ready!(self.conn.poll_open_send_stream(cx))?;
        Ok(stream.into()).into()
    }

    fn close(&mut self, code: h3::error::Code, _reason: &[u8]) {
        self.conn.close(
            code.value()
                .try_into()
                .unwrap_or_else(|_| VarInt::MAX.into()),
        );
    }
}

impl Clone for OpenStreams {
    fn clone(&self) -> Self {
        Self {
            conn: self.conn.clone(),
        }
    }
}

pub struct BidiStream<B>
where
    B: Buf,
{
    send: SendStream<B>,
    recv: RecvStream,
}

impl<B> quic::BidiStream<B> for BidiStream<B>
where
    B: Buf,
{
    type SendStream = SendStream<B>;
    type RecvStream = RecvStream;

    fn split(self) -> (Self::SendStream, Self::RecvStream) {
        (self.send, self.recv)
    }
}

impl<B> quic::RecvStream for BidiStream<B>
where
    B: Buf,
{
    type Buf = Bytes;
    type Error = ReadError;

    fn poll_data(
        &mut self,
        cx: &mut task::Context<'_>,
    ) -> Poll<Result<Option<Self::Buf>, Self::Error>> {
        self.recv.poll_data(cx)
    }

    fn stop_sending(&mut self, error_code: u64) {
        self.recv.stop_sending(error_code)
    }
}

impl<B> quic::SendStream<B> for BidiStream<B>
where
    B: Buf,
{
    type Error = SendStreamError;

    fn poll_ready(&mut self, cx: &mut task::Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.send.poll_ready(cx)
    }

    fn poll_finish(&mut self, cx: &mut task::Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.send.poll_finish(cx)
    }

    fn reset(&mut self, reset_code: u64) {
        self.send.reset(reset_code)
    }

    fn send_data<D: Into<WriteBuf<B>>>(&mut self, data: D) -> Result<(), Self::Error> {
        self.send.send_data(data)
    }

    fn id(&self) -> StreamId {
        self.send.id()
    }
}

impl<B> From<BidirectionalStream> for BidiStream<B>
where
    B: Buf,
{
    fn from(bidi: BidirectionalStream) -> Self {
        let (recv, send) = bidi.split();
        BidiStream {
            send: send.into(),
            recv: recv.into(),
        }
    }
}

pub struct RecvStream {
    stream: s2n_quic::stream::ReceiveStream,
}

impl RecvStream {
    fn new(stream: s2n_quic::stream::ReceiveStream) -> Self {
        Self { stream }
    }
}

impl quic::RecvStream for RecvStream {
    type Buf = Bytes;
    type Error = ReadError;

    fn poll_data(
        &mut self,
        cx: &mut task::Context<'_>,
    ) -> Poll<Result<Option<Self::Buf>, Self::Error>> {
        let buf = ready!(self.stream.poll_receive(cx))?;
        Ok(buf).into()
    }

    fn stop_sending(&mut self, error_code: u64) {
        let _ = self.stream.stop_sending(
            s2n_quic::application::Error::new(error_code)
                .expect("s2n-quic supports error codes up to 2^62-1"),
        );
    }
}

impl From<ReceiveStream> for RecvStream {
    fn from(recv: ReceiveStream) -> Self {
        RecvStream::new(recv)
    }
}

#[derive(Debug)]
pub struct ReadError(s2n_quic::stream::Error);

impl std::error::Error for ReadError {}

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

impl From<ReadError> for Arc<dyn Error> {
    fn from(e: ReadError) -> Self {
        Arc::new(e)
    }
}

impl From<s2n_quic::stream::Error> for ReadError {
    fn from(e: s2n_quic::stream::Error) -> Self {
        Self(e)
    }
}

impl Error for ReadError {
    fn is_timeout(&self) -> bool {
        matches!(
            self.0,
            s2n_quic::stream::Error::ConnectionError {
                error: s2n_quic::connection::Error::IdleTimerExpired { .. },
                ..
            }
        )
    }

    fn err_code(&self) -> Option<u64> {
        match self.0 {
            s2n_quic::stream::Error::ConnectionError {
                error: s2n_quic::connection::Error::Application { error, .. },
                ..
            } => Some(error.into()),
            s2n_quic::stream::Error::StreamReset { error, .. } => Some(error.into()),
            _ => None,
        }
    }
}

pub struct SendStream<B: Buf> {
    stream: s2n_quic::stream::SendStream,
    chunk: Option<Bytes>,
    buf: Option<WriteBuf<B>>, // TODO: Replace with buf: PhantomData<B>
                              //       after https://github.com/hyperium/h3/issues/78 is resolved
}

impl<B> SendStream<B>
where
    B: Buf,
{
    fn new(stream: s2n_quic::stream::SendStream) -> SendStream<B> {
        Self {
            stream,
            chunk: None,
            buf: Default::default(),
        }
    }
}

impl<B> quic::SendStream<B> for SendStream<B>
where
    B: Buf,
{
    type Error = SendStreamError;

    fn poll_ready(&mut self, cx: &mut task::Context<'_>) -> Poll<Result<(), Self::Error>> {
        loop {
            // try to flush the current chunk if we have one
            if let Some(chunk) = self.chunk.as_mut() {
                ready!(self.stream.poll_send(chunk, cx))?;

                // s2n-quic will take the whole chunk on send, even if it exceeds the limits
                debug_assert!(chunk.is_empty());
                self.chunk = None;
            }

            // try to take the next chunk from the WriteBuf
            if let Some(ref mut data) = self.buf {
                let len = data.chunk().len();

                // if the write buf is empty, then clear it and break
                if len == 0 {
                    self.buf = None;
                    break;
                }

                // copy the first chunk from WriteBuf and prepare it to flush
                let chunk = data.copy_to_bytes(len);
                self.chunk = Some(chunk);

                // loop back around to flush the chunk
                continue;
            }

            // if we didn't have either a chunk or WriteBuf, then we're ready
            break;
        }

        Poll::Ready(Ok(()))

        // TODO: Replace with following after https://github.com/hyperium/h3/issues/78 is resolved
        // self.available_bytes = ready!(self.stream.poll_send_ready(cx))?;
        // Poll::Ready(Ok(()))
    }

    fn send_data<D: Into<WriteBuf<B>>>(&mut self, data: D) -> Result<(), Self::Error> {
        if self.buf.is_some() {
            return Err(Self::Error::NotReady);
        }
        self.buf = Some(data.into());
        Ok(())

        // TODO: Replace with following after https://github.com/hyperium/h3/issues/78 is resolved
        // let mut data = data.into();
        // while self.available_bytes > 0 && data.has_remaining() {
        //     let len = data.chunk().len();
        //     let chunk = data.copy_to_bytes(len);
        //     self.stream.send_data(chunk)?;
        //     self.available_bytes = self.available_bytes.saturating_sub(len);
        // }
        // Ok(())
    }

    fn poll_finish(&mut self, cx: &mut task::Context<'_>) -> Poll<Result<(), Self::Error>> {
        // ensure all chunks are flushed to the QUIC stream before finishing
        ready!(self.poll_ready(cx))?;
        self.stream.finish()?;
        Ok(()).into()
    }

    fn reset(&mut self, reset_code: u64) {
        let _ = self
            .stream
            .reset(reset_code.try_into().unwrap_or_else(|_| VarInt::MAX.into()));
    }

    fn id(&self) -> StreamId {
        self.stream.id().try_into().expect("invalid stream id")
    }
}

impl<B> From<s2n_quic::stream::SendStream> for SendStream<B>
where
    B: Buf,
{
    fn from(send: s2n_quic::stream::SendStream) -> Self {
        SendStream::new(send)
    }
}

#[derive(Debug)]
pub enum SendStreamError {
    Write(s2n_quic::stream::Error),
    NotReady,
}

impl std::error::Error for SendStreamError {}

impl Display for SendStreamError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{self:?}")
    }
}

impl From<s2n_quic::stream::Error> for SendStreamError {
    fn from(e: s2n_quic::stream::Error) -> Self {
        Self::Write(e)
    }
}

impl Error for SendStreamError {
    fn is_timeout(&self) -> bool {
        matches!(
            self,
            Self::Write(s2n_quic::stream::Error::ConnectionError {
                error: s2n_quic::connection::Error::IdleTimerExpired { .. },
                ..
            })
        )
    }

    fn err_code(&self) -> Option<u64> {
        match self {
            Self::Write(s2n_quic::stream::Error::StreamReset { error, .. }) => {
                Some((*error).into())
            }
            Self::Write(s2n_quic::stream::Error::ConnectionError {
                error: s2n_quic::connection::Error::Application { error, .. },
                ..
            }) => Some((*error).into()),
            _ => None,
        }
    }
}

impl From<SendStreamError> for Arc<dyn Error> {
    fn from(e: SendStreamError) -> Self {
        Arc::new(e)
    }
}


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

use std::{fs::read_dir, io::Error, path::Path, process::Command};

fn main() -> Result<(), Error> {
    let env = Env::new();

    for feature in read_dir("features")? {
        let path = feature?.path();
        if let Some(name) = path.file_stem() {
            println!("cargo:rerun-if-changed={}", path.display());
            if env.check(&path)? {
                supports(name.to_str().expect("valid feature name"));
            }
        }
    }

    let is_miri = std::env::var("CARGO_CFG_MIRI").is_ok();

    match env.target_os.as_str() {
        "linux" => {
            supports("gso");
            supports("gro");
            supports("mtu_disc");
            supports("pktinfo");
            supports("tos");

            // miri doesn't support the way we detect syscall support so override it
            if is_miri {
                supports("socket_msg");
                supports("socket_mmsg");
            }
        }
        "macos" => {
            supports("pktinfo");
            supports("tos");

            // miri doesn't support the way we detect syscall support so override it
            if is_miri {
                supports("socket_msg");
            }
        }
        "android" => {
            supports("mtu_disc");
            supports("pktinfo");
            supports("tos");
        }
        _ => {
            // TODO others
        }
    }

    Ok(())
}

fn supports(name: &str) {
    println!("cargo:rustc-cfg=s2n_quic_platform_{name}");
}

struct Env {
    rustc: String,
    out_dir: String,
    target: String,
    target_os: String,
}

impl Env {
    fn new() -> Self {
        // See https://doc.rust-lang.org/cargo/reference/environment-variables.html#environment-variables-cargo-sets-for-build-scripts
        Self {
            rustc: env("RUSTC"),
            out_dir: env("OUT_DIR"),
            target: env("TARGET"),
            target_os: env("CARGO_CFG_TARGET_OS"),
        }
    }

    // Tries to compile the program and returns if it was successful
    fn check(&self, path: &Path) -> Result<bool, Error> {
        let mut command = Command::new(&self.rustc);

        command
            .arg("--out-dir")
            .arg(&self.out_dir)
            .arg("--target")
            .arg(&self.target)
            .arg("--crate-type")
            .arg("bin")
            .arg("--codegen")
            .arg("opt-level=0")
            .arg(path);

        for (key, _) in std::env::vars() {
            const CARGO_FEATURE: &str = "CARGO_FEATURE_";
            if key.starts_with(CARGO_FEATURE) {
                command.arg("--cfg").arg(format!(
                    "feature=\"{}\"",
                    key.trim_start_matches(CARGO_FEATURE)
                        .to_lowercase()
                        .replace('_', "-")
                ));
            }
        }

        Ok(command.spawn()?.wait()?.success())
    }
}

fn env(name: &str) -> String {
    println!("cargo:rerun-if-env-changed={name}");
    std::env::var(name)
        .unwrap_or_else(|_| panic!("build script missing {name:?} environment variable"))
}


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

//! Detects platform capability of:
//! * https://linux.die.net/man/2/sendmmsg
//! * https://linux.die.net/man/2/recvmmsg

fn main() {
    println!("sendmmsg {:?}", unsafe { SENDMMSG });
    println!("recvmmsg {:?}", unsafe { RECVMMSG });
}

/// Try to resolve the required references from the linker
///
/// The build will fail if they don't exist.
#[cfg(all(unix, feature = "std"))]
extern "C" {
    #[link_name = "sendmmsg"]
    static SENDMMSG: *const u8;
    #[link_name = "recvmmsg"]
    static RECVMMSG: *const u8;
}


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

//! Detects platform capability of:
//! * https://linux.die.net/man/2/sendmsg
//! * https://linux.die.net/man/2/recvmsg

fn main() {
    println!("sendmsg {:?}", unsafe { SENDMSG });
    println!("recvmsg {:?}", unsafe { RECVMSG });
}

/// Try to resolve the required references from the linker
///
/// The build will fail if they don't exist.
#[cfg(all(unix, feature = "std"))]
extern "C" {
    #[link_name = "sendmsg"]
    static SENDMSG: *const u8;
    #[link_name = "recvmsg"]
    static RECVMSG: *const u8;
}


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

pub mod gso;
pub use gso::Gso;


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

#[cfg(feature = "tokio")]
pub mod tokio;

#[cfg(any(test, feature = "io-testing"))]
pub mod testing;

#[cfg(feature = "turmoil")]
pub mod turmoil;

#[cfg(feature = "xdp")]
pub mod xdp;


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

//! This module contains abstractions around the platform on which the
//! stack is running

#![cfg_attr(not(any(test, feature = "std")), no_std)]

extern crate alloc;

mod features;
pub mod io;
mod message;
mod socket;
mod syscall;
#[doc(hidden)] // TODO remove this module: https://github.com/aws/s2n-quic/issues/1738
pub mod time;


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

use core::{alloc::Layout, ptr::NonNull};
use s2n_quic_core::{inet::datagram, io::tx, path};

#[cfg(any(s2n_quic_platform_socket_msg, s2n_quic_platform_socket_mmsg))]
mod cmsg;
#[cfg(s2n_quic_platform_socket_mmsg)]
pub mod mmsg;
#[cfg(s2n_quic_platform_socket_msg)]
pub mod msg;
pub mod simple;

pub mod default {
    cfg_if::cfg_if! {
        if #[cfg(s2n_quic_platform_socket_mmsg)] {
            pub use super::mmsg::*;
        } else if #[cfg(s2n_quic_platform_socket_msg)] {
            pub use super::msg::*;
        } else {
            pub use super::simple::*;
        }
    }
}

/// Tracks allocations of message ring buffer state
pub struct Storage {
    ptr: NonNull<u8>,
    layout: Layout,
}

impl Storage {
    #[inline]
    pub fn new(layout: Layout) -> Self {
        unsafe {
            let ptr = alloc::alloc::alloc_zeroed(layout);
            let ptr = NonNull::new(ptr).expect("could not allocate message storage");
            Self { layout, ptr }
        }
    }

    #[inline]
    pub(crate) fn as_ptr(&self) -> *mut u8 {
        self.ptr.as_ptr()
    }

    /// Asserts that the pointer is in bounds of the allocation
    #[inline]
    pub(crate) fn check_bounds<T: Sized>(&self, ptr: *mut T) {
        let start = self.as_ptr();
        let end = unsafe {
            // Safety: pointer is allocated with the self.layout
            start.add(self.layout.size())
        };
        let allocation_range = start..=end;
        let actual_end_ptr = ptr as *mut u8;
        debug_assert!(allocation_range.contains(&actual_end_ptr));
    }
}

impl Drop for Storage {
    fn drop(&mut self) {
        unsafe {
            // Safety: pointer was allocated with self.layout
            alloc::alloc::dealloc(self.as_ptr(), self.layout)
        }
    }
}

/// An abstract message that can be sent and received on a network
pub trait Message: 'static + Copy {
    type Handle: path::Handle;

    const SUPPORTS_GSO: bool;
    const SUPPORTS_ECN: bool;
    const SUPPORTS_FLOW_LABELS: bool;

    /// Allocates `entries` messages, each with `payload_len` bytes
    fn alloc(entries: u32, payload_len: u32, offset: usize) -> Storage;

    /// Returns the length of the payload
    fn payload_len(&self) -> usize;

    /// Sets the payload length for the message
    ///
    /// # Safety
    /// This method should only set the payload less than or
    /// equal to its initially allocated size.
    unsafe fn set_payload_len(&mut self, payload_len: usize);

    /// Validates that the `source` message can be replicated to `dest`.
    ///
    /// # Panics
    ///
    /// This panics when the messages cannot be replicated
    fn validate_replication(source: &Self, dest: &Self);

    /// Returns a mutable pointer for the message payload
    fn payload_ptr_mut(&mut self) -> *mut u8;

    /// Returns a mutable slice for the message payload
    #[inline]
    fn payload_mut(&mut self) -> &mut [u8] {
        unsafe { core::slice::from_raw_parts_mut(self.payload_ptr_mut(), self.payload_len()) }
    }

    /// Sets the segment size for the message payload
    fn set_segment_size(&mut self, _size: usize) {
        panic!("cannot use GSO on the current platform");
    }

    /// Resets the message for future use
    ///
    /// # Safety
    /// This method should only set the MTU to the original value
    unsafe fn reset(&mut self, mtu: usize);

    /// Reads the message as an RX packet
    fn rx_read(&mut self, local_address: &path::LocalAddress) -> Option<RxMessage<Self::Handle>>;

    /// Writes the message into the TX packet
    fn tx_write<M: tx::Message<Handle = Self::Handle>>(
        &mut self,
        message: M,
    ) -> Result<usize, tx::Error>;
}

pub struct RxMessage<'a, Handle: Copy> {
    /// The received header for the message
    pub header: datagram::Header<Handle>,
    /// The number of segments inside the message
    pub segment_size: usize,
    /// The full payload of the message
    pub payload: &'a mut [u8],
}

impl<'a, Handle: Copy> RxMessage<'a, Handle> {
    #[inline]
    pub fn for_each<F: FnMut(datagram::Header<Handle>, &mut [u8])>(self, mut on_packet: F) {
        // `chunks_mut` doesn't know what to do with zero-sized segments so return early
        if self.segment_size == 0 {
            return;
        }

        for segment in self.payload.chunks_mut(self.segment_size) {
            on_packet(self.header, segment);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bolero::check;

    #[test]
    #[cfg_attr(kani, kani::proof, kani::unwind(17), kani::solver(cadical))]
    fn rx_message_test() {
        let path = bolero::gen::<path::RemoteAddress>();
        let ecn = bolero::gen();
        let segment_size = bolero::gen();
        let max_payload_len = if cfg!(kani) { 16 } else { u16::MAX as usize };
        let payload_len = 0..=max_payload_len;

        check!()
            .with_generator((path, ecn, segment_size, payload_len))
            .cloned()
            .for_each(|(path, ecn, segment_size, payload_len)| {
                let mut payload = vec![0u8; payload_len];
                let rx_message = RxMessage {
                    header: datagram::Header { path, ecn },
                    segment_size,
                    payload: &mut payload,
                };

                rx_message.for_each(|header, segment| {
                    assert_eq!(header.path, path);
                    assert_eq!(header.ecn, ecn);
                    assert!(segment.len() <= payload_len);
                    assert!(segment.len() <= segment_size);
                })
            })
    }
}


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

pub mod io;
pub mod ring;
pub mod task;


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

// some platforms contain empty implementations so disable any warnings from those
#![allow(unused_variables, unused_macros, unused_mut, clippy::let_and_return)]

use core::ops::ControlFlow;
use socket2::{Domain, Protocol, Socket, Type};
use std::io;

/// Calls the given libc function and wraps the result in an `io::Result`.
macro_rules! libc {
    ($fn: ident ( $($arg: expr),* $(,)* ) ) => {{
        let res = unsafe { libc::$fn($($arg, )*) };
        if res < 0 {
            Err(std::io::Error::last_os_error())
        } else {
            Ok(res)
        }
    }};
}

#[cfg(s2n_quic_platform_socket_mmsg)]
pub mod mmsg;
#[cfg(s2n_quic_platform_socket_msg)]
pub mod msg;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[allow(dead_code)] // this isn't constructed on some platforms
pub enum SocketType {
    Blocking,
    NonBlocking,
}

pub trait SocketEvents {
    /// Called when `count` packets are completed
    ///
    /// If `Continue` is returned, the socket will assume the packet was acceptable and continue
    /// with the remaining packets. If `Break` is returned, the syscall stop looping and yield to
    /// the caller.
    fn on_complete(&mut self, count: usize) -> ControlFlow<(), ()>;

    /// Called when an error occurs on a socket
    ///
    /// If `Continue` is returned, the socket will discard the packet and continue
    /// with the remaining packets. If `Break` is returned, the syscall will assume the current
    /// packet can be retried and yield to the caller.
    fn on_error(&mut self, error: io::Error) -> ControlFlow<(), ()>;
}

#[cfg(unix)]
pub trait UnixMessage: crate::message::Message {
    fn send<E: SocketEvents>(fd: std::os::unix::io::RawFd, entries: &mut [Self], events: &mut E);
    fn recv<E: SocketEvents>(
        fd: std::os::unix::io::RawFd,
        ty: SocketType,
        entries: &mut [Self],
        events: &mut E,
    );
}

pub fn udp_socket(addr: std::net::SocketAddr) -> io::Result<Socket> {
    let domain = Domain::for_address(addr);
    let socket_type = Type::DGRAM;
    let protocol = Some(Protocol::UDP);

    let socket = Socket::new(domain, socket_type, protocol)?;

    // allow ipv4 to also connect - ignore the error if it fails
    let _ = socket.set_only_v6(false);

    socket.set_reuse_address(true)?;

    Ok(socket)
}

/// Creates a UDP socket bound to the provided address
pub fn bind_udp<A: std::net::ToSocketAddrs>(addr: A, reuse_port: bool) -> io::Result<Socket> {
    let addr = addr.to_socket_addrs()?.next().ok_or_else(|| {
        std::io::Error::new(
            io::ErrorKind::InvalidInput,
            "the provided bind address was empty",
        )
    })?;
    let socket = udp_socket(addr)?;

    #[cfg(unix)]
    socket.set_reuse_port(reuse_port)?;

    // mark the variable as "used" regardless of platform support
    let _ = reuse_port;

    socket.bind(&addr.into())?;

    Ok(socket)
}

/// Binds a socket to a specified interface by name
#[cfg(feature = "xdp")]
pub fn bind_to_interface<F: std::os::unix::io::AsRawFd>(
    socket: &F,
    ifname: &std::ffi::CStr,
) -> io::Result<()> {
    libc!(setsockopt(
        socket.as_raw_fd(),
        libc::SOL_SOCKET,
        libc::SO_BINDTODEVICE,
        ifname as *const _ as *const _,
        libc::IF_NAMESIZE as _
    ))?;
    Ok(())
}

/// Disables MTU discovery and fragmentation on the socket
pub fn configure_mtu_disc(tx_socket: &Socket) -> bool {
    let mut success = false;

    //= https://www.rfc-editor.org/rfc/rfc9000#section-14
    //# UDP datagrams MUST NOT be fragmented at the IP layer.

    //= https://www.rfc-editor.org/rfc/rfc9000#section-14
    //# In IPv4 [IPv4], the Don't Fragment (DF) bit MUST be set if possible, to
    //# prevent fragmentation on the path.

    //= https://www.rfc-editor.org/rfc/rfc8899#section-3
    //# In IPv4, a probe packet MUST be sent with the Don't
    //# Fragment (DF) bit set in the IP header and without network layer
    //# endpoint fragmentation.

    //= https://www.rfc-editor.org/rfc/rfc8899#section-4.5
    //# A PL implementing this specification MUST suspend network layer
    //# processing of outgoing packets that enforces a PMTU
    //# [RFC1191][RFC8201] for each flow utilizing DPLPMTUD and instead use
    //# DPLPMTUD to control the size of packets that are sent by a flow.
    #[cfg(s2n_quic_platform_mtu_disc)]
    {
        use std::os::unix::io::AsRawFd;

        // IP_PMTUDISC_PROBE setting will set the DF (Don't Fragment) flag
        // while also ignoring the Path MTU. This means packets will not
        // be fragmented, and the EMSGSIZE error will not be returned for
        // packets larger than the Path MTU according to the kernel.
        success |= libc!(setsockopt(
            tx_socket.as_raw_fd(),
            libc::IPPROTO_IP,
            libc::IP_MTU_DISCOVER,
            &libc::IP_PMTUDISC_PROBE as *const _ as _,
            core::mem::size_of_val(&libc::IP_PMTUDISC_PROBE) as _,
        ))
        .is_ok();

        success |= libc!(setsockopt(
            tx_socket.as_raw_fd(),
            libc::IPPROTO_IPV6,
            libc::IPV6_MTU_DISCOVER,
            &libc::IP_PMTUDISC_PROBE as *const _ as _,
            core::mem::size_of_val(&libc::IP_PMTUDISC_PROBE) as _,
        ))
        .is_ok();
    }

    success
}

/// Configures the socket to return TOS/ECN information as part of the ancillary data
pub fn configure_tos(rx_socket: &Socket) -> bool {
    let mut success = false;

    #[cfg(s2n_quic_platform_tos)]
    {
        use std::os::unix::io::AsRawFd;
        let enabled: libc::c_int = 1;

        success |= libc!(setsockopt(
            rx_socket.as_raw_fd(),
            libc::IPPROTO_IP,
            libc::IP_RECVTOS,
            &enabled as *const _ as _,
            core::mem::size_of_val(&enabled) as _,
        ))
        .is_ok();

        success |= libc!(setsockopt(
            rx_socket.as_raw_fd(),
            libc::IPPROTO_IPV6,
            libc::IPV6_RECVTCLASS,
            &enabled as *const _ as _,
            core::mem::size_of_val(&enabled) as _,
        ))
        .is_ok()
    }

    success
}

/// Configures the socket to return local address and interface information as part of the
/// ancillary data
pub fn configure_pktinfo(rx_socket: &Socket) -> bool {
    let mut success = false;

    // Set up the RX socket to pass information about the local address and interface
    #[cfg(s2n_quic_platform_pktinfo)]
    {
        use std::os::unix::io::AsRawFd;
        let enabled: libc::c_int = 1;

        success |= libc!(setsockopt(
            rx_socket.as_raw_fd(),
            libc::IPPROTO_IP,
            libc::IP_PKTINFO,
            &enabled as *const _ as _,
            core::mem::size_of_val(&enabled) as _,
        ))
        .is_ok();

        success |= libc!(setsockopt(
            rx_socket.as_raw_fd(),
            libc::IPPROTO_IPV6,
            libc::IPV6_RECVPKTINFO,
            &enabled as *const _ as _,
            core::mem::size_of_val(&enabled) as _,
        ))
        .is_ok();
    }

    success
}

pub fn configure_gro(rx_socket: &Socket) -> bool {
    let mut success = false;

    #[cfg(s2n_quic_platform_gro)]
    {
        use std::os::unix::io::AsRawFd;
        let enabled: libc::c_int = 1;

        success |= libc!(setsockopt(
            rx_socket.as_raw_fd(),
            libc::SOL_UDP,
            libc::UDP_GRO,
            &enabled as *const _ as _,
            core::mem::size_of_val(&enabled) as _
        ))
        .is_ok();
    }

    success
}


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

use core::{
    convert::{TryFrom, TryInto},
    fmt,
    fmt::{Display, Formatter},
    num::NonZeroUsize,
};

#[derive(Clone, Copy, Debug)]
pub struct MaxSegments(NonZeroUsize);

impl MaxSegments {
    pub(crate) const MAX: Self = gso_impl::MAX;
    pub(crate) const DEFAULT: Self = gso_impl::DEFAULT;
}

impl Default for MaxSegments {
    fn default() -> Self {
        MaxSegments::DEFAULT
    }
}

impl TryFrom<usize> for MaxSegments {
    type Error = MaxSegmentsError;

    fn try_from(value: usize) -> Result<Self, Self::Error> {
        if !(1..=Self::MAX.into()).contains(&value) {
            return Err(MaxSegmentsError);
        }

        Ok(MaxSegments(value.try_into().expect(
            "Value must be greater than zero according to the check above",
        )))
    }
}

#[derive(Debug)]
pub struct MaxSegmentsError;

impl Display for MaxSegmentsError {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "MaxSegments must be a non-zero value less than or equal to {}",
            MaxSegments::MAX.0.get()
        )
    }
}

impl From<MaxSegments> for usize {
    #[inline]
    fn from(value: MaxSegments) -> Self {
        value.0.get()
    }
}

#[cfg(s2n_quic_platform_gso)]
mod gso_enabled {
    use super::*;
    use std::sync::{
        atomic::{AtomicUsize, Ordering},
        Arc,
    };

    // This value represents the Maximum value MaxSegments can be set to, i.e. a Max of a Max. The
    // value comes from the Linux kernel:
    //
    // https://github.com/torvalds/linux/blob/e9f1cbc0c4114880090c7a578117d3b9cf184ad4/tools/testing/selftests/net/udpgso.c#L37
    // ```
    // #define UDP_MAX_SEGMENTS	(1 << 6UL)
    // ```
    pub const MAX: MaxSegments = MaxSegments(unsafe { NonZeroUsize::new_unchecked(1 << 6) });

    // The packet pacer enforces a burst limit of 10 packets, so generally there is no benefit to
    // exceeding that value for GSO segments. However, in low RTT/high bandwidth networks the pacing
    // interval may drop below the timer granularity, resulting in `MAX_BURST_PACKETS` being
    // exceeded. In such networks, setting a MaxSegments size higher than the default may have a
    // positive effect on efficiency.
    //= https://www.rfc-editor.org/rfc/rfc9002#section-7.7
    //# Senders SHOULD limit bursts to the initial congestion window
    pub const DEFAULT: MaxSegments = MaxSegments(unsafe {
        NonZeroUsize::new_unchecked(s2n_quic_core::recovery::MAX_BURST_PACKETS as usize)
    });

    #[derive(Clone, Debug)]
    pub struct Gso(Arc<AtomicUsize>);

    impl Default for Gso {
        fn default() -> Self {
            MaxSegments::DEFAULT.into()
        }
    }

    impl Gso {
        #[inline]
        pub fn max_segments(&self) -> usize {
            self.0.load(Ordering::Relaxed)
        }

        #[inline]
        pub fn disable(&self) {
            self.0.store(1, Ordering::Relaxed);
        }
    }

    impl From<MaxSegments> for Gso {
        #[inline]
        fn from(segments: MaxSegments) -> Self {
            Self(Arc::new(AtomicUsize::new(segments.0.into())))
        }
    }
}

#[cfg(any(not(s2n_quic_platform_gso), test))]
mod gso_disabled {
    #![cfg_attr(test, allow(dead_code))]

    use super::*;

    pub const MAX: MaxSegments = MaxSegments(unsafe { NonZeroUsize::new_unchecked(1) });
    pub const DEFAULT: MaxSegments = MAX;

    #[derive(Clone, Default, Debug)]
    pub struct Gso(());

    impl Gso {
        #[inline]
        pub fn max_segments(&self) -> usize {
            1
        }

        #[inline]
        #[allow(dead_code)] // this may or may not be used on certain platforms
        pub fn disable(&self) {
            // it's already disabled
        }
    }

    impl From<MaxSegments> for Gso {
        #[inline]
        fn from(_segments: MaxSegments) -> Self {
            Self(())
        }
    }
}

mod gso_impl {
    #[cfg(not(s2n_quic_platform_gso))]
    pub use super::gso_disabled::*;
    #[cfg(s2n_quic_platform_gso)]
    pub use super::gso_enabled::*;
}

pub use gso_impl::Gso;


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

use bach::time::scheduler;
use core::task::Poll;
use s2n_quic_core::{
    endpoint::Endpoint, inet::SocketAddress, io::event_loop::EventLoop, path::MaxMtu,
};

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

mod model;
pub mod network;
pub mod time;

pub use model::{Model, TxRecorder};
pub use network::{Network, PathHandle};
pub use time::now;

pub use bach::task::{self, primary, spawn};

pub mod rand {
    pub use ::bach::rand::*;

    #[derive(Clone, Copy, Default)]
    pub struct Havoc;

    impl s2n_quic_core::havoc::Random for Havoc {
        #[inline]
        fn fill(&mut self, bytes: &mut [u8]) {
            fill_bytes(bytes);
        }

        #[inline]
        fn gen_bool(&mut self) -> bool {
            gen()
        }

        #[inline]
        fn shuffle(&mut self, bytes: &mut [u8]) {
            shuffle(bytes);
        }

        #[inline]
        fn gen_range(&mut self, range: core::ops::Range<u64>) -> u64 {
            gen_range(range)
        }
    }
}

pub mod executor {
    pub use bach::executor::{Handle, JoinHandle};
}

pub struct Executor<N: Network> {
    executor: bach::executor::Executor<Env<N>>,
    handle: Handle,
}

impl<N: Network> Executor<N> {
    pub fn new(network: N, seed: u64) -> Self {
        let mut executor = bach::executor::Executor::new(|handle| Env {
            handle: handle.clone(),
            time: scheduler::Scheduler::new(),
            rand: bach::rand::Scope::new(seed),
            buffers: network::Buffers::default(),
            network,
            stalled_iterations: 0,
        });

        let handle = Handle {
            executor: executor.handle().clone(),
            buffers: executor.environment().buffers.clone(),
        };

        Self { executor, handle }
    }

    pub fn handle(&self) -> &Handle {
        &self.handle
    }

    pub fn enter<F: FnOnce() -> O, O>(&mut self, f: F) -> O {
        self.executor.environment().enter(f)
    }

    pub fn run(&mut self) {
        self.executor.block_on_primary();
    }

    pub fn close(&mut self) {
        // close the environment, which notifies all of the tasks that we're shutting down
        self.executor.environment().close(|| {});
        while self.executor.macrostep() > 0 {}

        // then close the actual executor
        self.executor.close()
    }
}

impl<N: Network> Drop for Executor<N> {
    fn drop(&mut self) {
        self.close();
    }
}

struct Env<N> {
    handle: bach::executor::Handle,
    time: scheduler::Scheduler,
    rand: bach::rand::Scope,
    buffers: network::Buffers,
    network: N,
    stalled_iterations: usize,
}

impl<N> Env<N> {
    fn enter<F: FnOnce() -> O, O>(&self, f: F) -> O {
        self.handle.enter(|| self.time.enter(|| self.rand.enter(f)))
    }

    fn close<F: FnOnce()>(&mut self, f: F) {
        let handle = &mut self.handle;
        let rand = &mut self.rand;
        let time = &mut self.time;
        let buffers = &mut self.buffers;
        handle.enter(|| {
            rand.enter(|| {
                time.close();
                time.enter(|| {
                    buffers.close();
                    f();
                });
            })
        })
    }
}

impl<N: Network> bach::executor::Environment for Env<N> {
    fn run<Tasks, F>(&mut self, tasks: Tasks) -> Poll<()>
    where
        Tasks: Iterator<Item = F> + Send,
        F: 'static + FnOnce() -> Poll<()> + Send,
    {
        let mut is_ready = true;

        let Self {
            handle,
            time,
            rand,
            buffers,
            network,
            ..
        } = self;

        handle.enter(|| {
            time.enter(|| {
                rand.enter(|| {
                    for task in tasks {
                        is_ready &= task().is_ready();
                    }
                    network.execute(buffers);
                })
            })
        });

        if is_ready {
            Poll::Ready(())
        } else {
            Poll::Pending
        }
    }

    fn on_macrostep(&mut self, count: usize) {
        // only advance time after a stall
        if count > 0 {
            self.stalled_iterations = 0;
            return;
        }

        self.stalled_iterations += 1;

        // A stalled iteration is a macrostep that didn't actually execute any tasks.
        //
        // The idea with limiting it prevents the runtime from looping endlessly and not
        // actually doing any work. The value of 100 was chosen somewhat arbitrarily as a high
        // enough number that we won't get false positives but low enough that the number of
        // loops stays within reasonable ranges.
        if self.stalled_iterations > 100 {
            panic!("the runtime stalled after 100 iterations");
        }

        while let Some(time) = self.time.advance() {
            let _ = time;
            if self.time.wake() > 0 {
                // if a task has woken, then reset the stall count
                self.stalled_iterations = 0;
                break;
            }
        }
    }

    fn close<F>(&mut self, close: F)
    where
        F: 'static + FnOnce() + Send,
    {
        Self::close(self, close)
    }
}

#[derive(Clone)]
pub struct Handle {
    executor: executor::Handle,
    buffers: network::Buffers,
}

impl Handle {
    pub fn builder(&self) -> Builder {
        Builder {
            handle: self.clone(),
            address: None,
            max_mtu: MaxMtu::default(),
        }
    }
}

pub struct Builder {
    handle: Handle,
    address: Option<SocketAddress>,
    max_mtu: MaxMtu,
}

impl Builder {
    pub fn build(self) -> Result<Io> {
        Ok(Io { builder: self })
    }

    pub fn with_max_mtu(mut self, max_mtu: u16) -> Self {
        self.max_mtu = max_mtu.try_into().unwrap();
        self
    }
}

pub struct Io {
    builder: Builder,
}

impl Io {
    pub fn start<E: Endpoint<PathHandle = network::PathHandle>>(
        self,
        mut endpoint: E,
    ) -> Result<(executor::JoinHandle<()>, SocketAddress)> {
        let Builder {
            handle: Handle { executor, buffers },
            address,
            max_mtu,
        } = self.builder;
        endpoint.set_max_mtu(max_mtu);

        let handle = address.unwrap_or_else(|| buffers.generate_addr());

        let (tx, rx) = buffers.register(handle);

        let clock = time::Clock::default();

        let event_loop = EventLoop {
            endpoint,
            clock,
            tx,
            rx,
        };
        let join = executor.spawn(event_loop.start());
        Ok((join, handle))
    }
}