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

use crate::{
    ack::{
        ack_eliciting_transmission::{AckElicitingTransmission, AckElicitingTransmissionSet},
        ack_ranges::{AckRanges, AckRangesError},
        ack_transmission_state::AckTransmissionState,
    },
    contexts::WriteContext,
    processed_packet::ProcessedPacket,
    transmission,
};
use s2n_quic_core::{
    ack,
    counter::{Counter, Saturating},
    event::{
        self,
        builder::{AckAction, AckProcessed},
        IntoEvent as _,
    },
    frame::{ack::EcnCounts, Ack, Ping},
    packet::number::{PacketNumber, PacketNumberSpace},
    time::{timer, Timer, Timestamp},
    varint::VarInt,
};

//= https://www.rfc-editor.org/rfc/rfc9000#section-13.2
//# Endpoints acknowledge all packets they receive and process.  However,
//# only ack-eliciting packets cause an ACK frame to be sent within the
//# maximum ack delay.  Packets that are not ack-eliciting are only
//# acknowledged when an ACK frame is sent for other reasons.
//#
//# When sending a packet for any reason, an endpoint SHOULD attempt to
//# include an ACK frame if one has not been sent recently.  Doing so
//# helps with timely loss detection at the peer.
//#
//# In general, frequent feedback from a receiver improves loss and
//# congestion response, but this has to be balanced against excessive
//# load generated by a receiver that sends an ACK frame in response to
//# every ack-eliciting packet.  The guidance offered below seeks to
//# strike this balance.

#[derive(Clone, Debug)]
pub struct AckManager {
    /// Time at which the AckManager will wake and transmit an ACK
    ack_delay_timer: Timer,

    /// Used to track the ACK-eliciting transmissions sent from the AckManager
    ack_eliciting_transmissions: AckElicitingTransmissionSet,

    /// All of the processed AckRanges that need to be ACKed
    pub(super) ack_ranges: AckRanges,

    /// Locally configured AckSettings
    pub ack_settings: ack::Settings,

    /// The largest packet number that we've acked - used for pn decoding
    largest_received_packet_number_acked: PacketNumber,

    /// The time at which we received the largest pn
    largest_received_packet_number_at: Option<Timestamp>,

    /// The number of processed packets since transmission
    processed_packets_since_transmission: Counter<u8, Saturating>,

    /// The number of transmissions since the last ACK-eliciting packet was sent
    transmissions_since_elicitation: Counter<u8, Saturating>,

    /// Used to transition through transmission/retransmission states
    transmission_state: AckTransmissionState,

    /// Explicit Congestion Notification counts from processed packets
    ecn_counts: EcnCounts,
}

impl AckManager {
    pub fn new(packet_space: PacketNumberSpace, ack_settings: ack::Settings) -> Self {
        Self {
            ack_delay_timer: Timer::default(),
            ack_eliciting_transmissions: AckElicitingTransmissionSet::default(),
            ack_settings,
            ack_ranges: AckRanges::new(ack_settings.ack_ranges_limit as usize),
            largest_received_packet_number_acked: packet_space
                .new_packet_number(VarInt::from_u8(0)),
            largest_received_packet_number_at: None,
            processed_packets_since_transmission: Counter::new(0),
            transmissions_since_elicitation: Counter::new(0),
            transmission_state: AckTransmissionState::default(),
            ecn_counts: EcnCounts::default(),
        }
    }

    /// Called when an outgoing packet is being assembled
    pub fn on_transmit<W: WriteContext>(&mut self, context: &mut W) -> bool {
        let constraint = context.transmission_constraint();
        let mode = context.transmission_mode();
        let has_ranges = !self.ack_ranges.is_empty();
        let should_transmit = self
            .transmission_state
            .should_transmit(constraint, mode, has_ranges);

        if !should_transmit {
            return false;
        }

        let ack_delay = self.ack_delay(context.current_time());

        //= https://www.rfc-editor.org/rfc/rfc9000#section-13.4.1
        //# Even if an endpoint does not set an ECT field on packets it sends,
        //# the endpoint MUST provide feedback about ECN markings it receives, if
        //# these are accessible.
        context
            .write_ack_frame(&Ack {
                ack_delay,
                ack_ranges: &self.ack_ranges,
                ecn_counts: self.ecn_counts.as_option(),
            })
            .is_some()
    }

    /// Called after an outgoing packet is assembled and `on_transmit` returned `true`
    pub fn on_transmit_complete<W: WriteContext>(&mut self, context: &mut W) {
        debug_assert!(
            self.transmission_state.should_transmit(
                context.transmission_constraint(),
                context.transmission_mode(),
                !self.ack_ranges.is_empty()
            ),
            "`on_transmit_complete` was called when `should_transmit` is false"
        );

        // if we transmitted something no need to wake up again to transmit the same thing
        self.ack_delay_timer.cancel();

        let mut is_ack_eliciting = context.ack_elicitation().is_ack_eliciting();

        if !is_ack_eliciting {
            // check the timer and make sure we can still write a Ping frame before removing it
            // We send a ping even when constrained to retransmissions only, as a fast
            // retransmission that is not ack eliciting will not help us recover faster.
            if (context.transmission_constraint().can_transmit()
                || context.transmission_constraint().can_retransmit())
                && self.transmissions_since_elicitation
                    >= self.ack_settings.ack_elicitation_interval
                && context.write_frame(&Ping).is_some()
            {
                is_ack_eliciting = true;
            } else {
                self.transmissions_since_elicitation += 1;
            }
        }

        self.largest_received_packet_number_acked = self
            .ack_ranges
            .max_value()
            .expect("transmission_state should be Disabled while ack_ranges is empty");

        if is_ack_eliciting {
            // reset the counter
            self.transmissions_since_elicitation = Counter::new(0);

            //= https://www.rfc-editor.org/rfc/rfc9000#section-13.2.4
            //# When a packet containing an ACK frame is sent, the Largest
            //# Acknowledged field in that frame can be saved.
            self.ack_eliciting_transmissions
                .on_transmit(AckElicitingTransmission {
                    sent_in_packet: context.packet_number(),
                    largest_received_packet_number_acked: self.largest_received_packet_number_acked,
                });
        }

        // record a transmission
        self.transmission_state
            .on_transmit(!self.ack_ranges.is_empty());

        // reset the number of packets since transmission
        self.processed_packets_since_transmission = Counter::new(0);
    }

    /// Called when a set of packets was acknowledged
    pub fn on_packet_ack<A: ack::Set>(&mut self, _timestamp: Timestamp, ack_set: &A) {
        if let Some(ack_range) = self.ack_eliciting_transmissions.on_update(ack_set) {
            self.ack_ranges
                .remove(ack_range)
                .expect("The range should always shrink the interval length");

            // `self.transmission_state` will be automatically notified in `on_processed_packet`
            // so wait for that instead
        }
    }

    /// Called when a set of packets was reported lost
    pub fn on_packet_loss<A: ack::Set>(&mut self, ack_set: &A) {
        if self
            .ack_eliciting_transmissions
            .on_update(ack_set)
            .is_some()
        {
            // transition to active mode when packet is lost
            self.transmission_state.on_update(&self.ack_ranges);
            self.transmission_state.activate();
        }
    }

    /// Called after an RX packet has been processed
    pub fn on_processed_packet<Pub: event::ConnectionPublisher>(
        &mut self,
        processed_packet: &ProcessedPacket,
        path: event::builder::Path,
        publisher: &mut Pub,
    ) {
        let packet_number = processed_packet.packet_number;
        let now = processed_packet.datagram.timestamp;

        // perform some checks before inserting into the ack_ranges
        let (is_ordered, is_largest) = self
            .ack_ranges
            .max_value()
            .and_then(|max_value| {
                // check to see if the packet number is the next one in the sequence
                let is_ordered = packet_number == max_value.next()?;

                // check to see if the packet number is the largest we've seen
                let is_largest = packet_number > max_value;

                Some((is_ordered, is_largest))
            })
            .unwrap_or((true, true));

        // This will fail if `packet_number` is less than `ack_ranges.min_value()`
        // and `ack_ranges` is at capacity.
        //
        // Most likely, this packet is very old and the contents have already
        // been retransmitted by the peer.
        if let Err(err) = self.ack_ranges.insert_packet_number(packet_number) {
            match err {
                AckRangesError::RangeInsertionFailed { min, max }
                | AckRangesError::LowestRangeDropped { min, max } => {
                    let start = self
                        .ack_ranges
                        .min_value()
                        .expect("should be non empty")
                        .into_event();
                    let end = self
                        .ack_ranges
                        .max_value()
                        .expect("should be non empty")
                        .into_event();

                    publisher.on_rx_ack_range_dropped(event::builder::RxAckRangeDropped {
                        path,
                        packet_number_range: min.into_event()..=max.into_event(),
                        capacity: self.ack_ranges.interval_len(),
                        stored_range: start..=end,
                    });

                    publisher.on_ack_processed(AckProcessed {
                        action: AckAction::RxAckRangeDropped {
                            packet_number_range: min.into_event()..=max.into_event(),
                            capacity: self.ack_ranges.interval_len(),
                            stored_range: start..=end,
                        },
                        path,
                    });
                }
            };
        }

        //= https://www.rfc-editor.org/rfc/rfc9000#section-13.4.1
        //# ECN counts are only incremented when QUIC packets from the received
        //# IP packet are processed.  As such, duplicate QUIC packets are not
        //# processed and do not increase ECN counts; see Section 21.10 for
        //# relevant security concerns.
        self.ecn_counts.increment(processed_packet.datagram.ecn);

        // Notify the state that the ack_ranges have changed
        self.transmission_state.on_update(&self.ack_ranges);
        self.processed_packets_since_transmission += 1;

        //= https://www.rfc-editor.org/rfc/rfc9000#section-13.2.5
        //# An endpoint measures the delays intentionally introduced between the
        //# time the packet with the largest packet number is received and the
        //# time an acknowledgment is sent.  The endpoint encodes this
        //# acknowledgement delay in the ACK Delay field of an ACK frame; see
        //# Section 19.3.  This allows the receiver of the ACK frame to adjust
        //# for any intentional delays, which is important for getting a better
        //# estimate of the path RTT when acknowledgments are delayed.
        if is_largest {
            self.largest_received_packet_number_at = Some(now);
        }

        if processed_packet.is_ack_eliciting() {
            let mut should_activate = false;

            //= https://www.rfc-editor.org/rfc/rfc9000#section-13.2.1
            //# In order to assist loss detection at the sender, an endpoint SHOULD
            //# generate and send an ACK frame without delay when it receives an ack-
            //# eliciting packet either:
            //#
            //# *  when the received packet has a packet number less than another
            //#    ack-eliciting packet that has been received, or

            should_activate |= !is_largest;

            //= https://www.rfc-editor.org/rfc/rfc9000#section-13.2.1
            //# *  when the packet has a packet number larger than the highest-
            //#    numbered ack-eliciting packet that has been received and there are
            //#    missing packets between that packet and this packet.

            should_activate |= !is_ordered;

            //= https://www.rfc-editor.org/rfc/rfc9000#section-13.2.1
            //# Similarly, packets marked with the ECN Congestion Experienced (CE)
            //# codepoint in the IP header SHOULD be acknowledged immediately, to
            //# reduce the peer's response time to congestion events.
            should_activate |= processed_packet.datagram.ecn.congestion_experienced();

            // TODO update to draft link after published
            // https://github.com/quicwg/base-drafts/pull/3623
            // An ACK frame SHOULD be generated for at least every 10th ack-eliciting packet

            // TODO support delayed ack proposal
            // https://tools.ietf.org/html/draft-iyengar-quic-delayed-ack-00
            let packet_tolerance = 10;

            should_activate |= self.processed_packets_since_transmission >= packet_tolerance;

            //= https://www.rfc-editor.org/rfc/rfc9000#section-9.3.3
            //# An endpoint that receives a PATH_CHALLENGE on an active path SHOULD
            //# send a non-probing packet in response.
            should_activate |= processed_packet.path_challenge_on_active_path;

            if should_activate {
                self.transmission_state.activate();
            } else if !self.ack_delay_timer.is_armed() {
                //= https://www.rfc-editor.org/rfc/rfc9000#section-13.2
                //# Endpoints acknowledge all packets they receive and process.  However,
                //# only ack-eliciting packets cause an ACK frame to be sent within the
                //# maximum ack delay.  Packets that are not ack-eliciting are only
                //# acknowledged when an ACK frame is sent for other reasons.
                self.ack_delay_timer
                    .set(now + self.ack_settings.max_ack_delay)
            }
        }

        // To save on timer churn, check to see if we've already expired since the
        // last time we sent an ACK frame
        if self.ack_delay_timer.poll_expiration(now).is_ready() {
            self.transmission_state.activate();
        }
    }

    /// Called when the connection timer expired
    pub fn on_timeout(&mut self, timestamp: Timestamp) {
        // NOTE: ack_elicitation_timer is not actively polled

        if self.ack_delay_timer.poll_expiration(timestamp).is_ready() {
            // transition to active transmission when we exceed the ack_delay
            self.transmission_state.activate();
        }
    }

    /// Returns the largest received packet number that has been ACKed at least once
    pub fn largest_received_packet_number_acked(&self) -> PacketNumber {
        self.largest_received_packet_number_acked
    }

    /// Computes the ack_delay field for the current state
    fn ack_delay(&self, now: Timestamp) -> VarInt {
        let ack_delay = self
            .largest_received_packet_number_at
            .map(|prev| now.saturating_duration_since(prev))
            .unwrap_or_default();
        self.ack_settings.encode_ack_delay(ack_delay)
    }
}

impl timer::Provider for AckManager {
    #[inline]
    fn timers<Q: timer::Query>(&self, query: &mut Q) -> timer::Result {
        // NOTE: ack_elicitation_timer is not actively polled

        self.ack_delay_timer.timers(query)?;

        Ok(())
    }
}

impl transmission::interest::Provider for AckManager {
    #[inline]
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        self.transmission_state.transmission_interest(query)
    }
}

#[cfg(test)]
mod tests {
    use super::{super::tests::*, *};
    use crate::{
        contexts::testing::{MockWriteContext, OutgoingFrameBuffer},
        path::{path_event, testing::helper_path_server},
    };
    use core::{
        iter::{empty, once},
        time::Duration,
    };
    use insta::assert_debug_snapshot;
    use s2n_quic_core::{
        ack, connection, endpoint,
        event::testing::Publisher,
        frame::{ack_elicitation::AckElicitation, ping, Frame},
        inet::{DatagramInfo, ExplicitCongestionNotification},
        path,
        time::{clock::testing as time, Clock, NoopClock},
    };

    #[test]
    fn activate() {
        // Setup:
        let mut manager =
            AckManager::new(PacketNumberSpace::ApplicationData, ack::Settings::default());

        let pn = PacketNumberSpace::ApplicationData.new_packet_number(VarInt::from_u8(1));
        let datagram = DatagramInfo {
            ecn: Default::default(),
            payload_len: 1200,
            timestamp: NoopClock {}.get_time(),
            destination_connection_id: connection::LocalId::TEST_ID,
            source_connection_id: None,
        };
        let mut processed_packet = ProcessedPacket::new(pn, &datagram);
        processed_packet.path_challenge_on_active_path = true;
        processed_packet.ack_elicitation = AckElicitation::Eliciting;

        assert!(!manager.transmission_state.is_active());

        // Trigger:
        let path = helper_path_server();
        let path_id = path::Id::test_id();
        manager.on_processed_packet(
            &processed_packet,
            path_event!(path, path_id),
            &mut Publisher::snapshot(),
        );

        // Expectation:
        assert!(manager.transmission_state.is_active());
    }

    #[test]
    fn ecn_counts() {
        // Setup:
        let mut manager =
            AckManager::new(PacketNumberSpace::ApplicationData, ack::Settings::default());

        assert_eq!(0, manager.ecn_counts.ect_0_count.as_u64());
        assert_eq!(0, manager.ecn_counts.ect_1_count.as_u64());
        assert_eq!(0, manager.ecn_counts.ce_count.as_u64());

        // Process a packet in an Ect0 datagram
        let pn = PacketNumberSpace::ApplicationData.new_packet_number(VarInt::from_u8(1));
        let datagram = helper_datagram_info(ExplicitCongestionNotification::Ect0);
        let path = helper_path_server();
        let path_id = path::Id::test_id();
        let mut publisher = Publisher::snapshot();
        manager.on_processed_packet(
            &ProcessedPacket::new(pn, &datagram),
            path_event!(path, path_id),
            &mut publisher,
        );

        assert_eq!(1, manager.ecn_counts.ect_0_count.as_u64());
        assert_eq!(0, manager.ecn_counts.ect_1_count.as_u64());
        assert_eq!(0, manager.ecn_counts.ce_count.as_u64());

        // Process a couple packets in an Ect1 datagram
        let pn1 = PacketNumberSpace::ApplicationData.new_packet_number(VarInt::from_u8(2));
        let pn2 = PacketNumberSpace::ApplicationData.new_packet_number(VarInt::from_u8(3));
        let datagram = helper_datagram_info(ExplicitCongestionNotification::Ect1);

        manager.on_processed_packet(
            &ProcessedPacket::new(pn1, &datagram),
            path_event!(path, path_id),
            &mut publisher,
        );
        manager.on_processed_packet(
            &ProcessedPacket::new(pn2, &datagram),
            path_event!(path, path_id),
            &mut publisher,
        );

        assert_eq!(1, manager.ecn_counts.ect_0_count.as_u64());
        assert_eq!(2, manager.ecn_counts.ect_1_count.as_u64());
        assert_eq!(0, manager.ecn_counts.ce_count.as_u64());

        // Process a packet in an Ce datagram
        let pn = PacketNumberSpace::ApplicationData.new_packet_number(VarInt::from_u8(4));
        let datagram = helper_datagram_info(ExplicitCongestionNotification::Ce);
        manager.on_processed_packet(
            &ProcessedPacket::new(pn, &datagram),
            path_event!(path, path_id),
            &mut publisher,
        );

        assert_eq!(1, manager.ecn_counts.ect_0_count.as_u64());
        assert_eq!(2, manager.ecn_counts.ect_1_count.as_u64());
        assert_eq!(1, manager.ecn_counts.ce_count.as_u64());

        // Process a packet in a NotEct datagram
        let pn = PacketNumberSpace::ApplicationData.new_packet_number(VarInt::from_u8(5));
        let datagram = helper_datagram_info(ExplicitCongestionNotification::NotEct);
        manager.on_processed_packet(
            &ProcessedPacket::new(pn, &datagram),
            path_event!(path, path_id),
            &mut publisher,
        );

        assert_eq!(1, manager.ecn_counts.ect_0_count.as_u64());
        assert_eq!(2, manager.ecn_counts.ect_1_count.as_u64());
        assert_eq!(1, manager.ecn_counts.ce_count.as_u64());
    }

    /// Helper function to construct `DatagramInfo` with the given `ExplicitCongestionNotification`
    fn helper_datagram_info(ecn: ExplicitCongestionNotification) -> DatagramInfo {
        DatagramInfo {
            ecn,
            payload_len: 1200,
            timestamp: NoopClock {}.get_time(),
            destination_connection_id: connection::LocalId::TEST_ID,
            source_connection_id: None,
        }
    }

    #[test]
    fn on_transmit_complete_transmission_constrained() {
        let mut manager =
            AckManager::new(PacketNumberSpace::ApplicationData, ack::Settings::default());
        let mut frame_buffer = OutgoingFrameBuffer::new();
        let mut write_context = MockWriteContext::new(
            time::now(),
            &mut frame_buffer,
            transmission::Constraint::None,
            transmission::Mode::Normal,
            endpoint::Type::Server,
        );

        manager.ack_ranges = AckRanges::default();
        assert!(manager
            .ack_ranges
            .insert_packet_number(
                PacketNumberSpace::ApplicationData.new_packet_number(VarInt::from_u8(1)),
            )
            .is_ok());
        manager.transmission_state = AckTransmissionState::Active { retransmissions: 0 };
        manager.transmissions_since_elicitation =
            Counter::new(ack::Settings::EARLY.ack_elicitation_interval);

        manager.on_transmit_complete(&mut write_context);

        assert_eq!(
            write_context
                .frame_buffer
                .pop_front()
                .expect("Frame is written")
                .as_frame(),
            Frame::Ping(ping::Ping),
            "Ping should be written when transmission is not constrained"
        );

        manager.transmission_state = AckTransmissionState::Active { retransmissions: 0 };
        manager.transmissions_since_elicitation =
            Counter::new(ack::Settings::EARLY.ack_elicitation_interval);
        write_context.frame_buffer.clear();
        write_context.transmission_constraint = transmission::Constraint::CongestionLimited;

        manager.on_transmit_complete(&mut write_context);
        assert!(
            write_context.frame_buffer.is_empty(),
            "Ping should not be written when CongestionLimited"
        );

        manager.transmission_state = AckTransmissionState::Active { retransmissions: 0 };
        manager.transmissions_since_elicitation =
            Counter::new(ack::Settings::EARLY.ack_elicitation_interval);
        write_context.frame_buffer.clear();
        write_context.transmission_constraint = transmission::Constraint::RetransmissionOnly;

        manager.on_transmit_complete(&mut write_context);
        assert_eq!(
            write_context
                .frame_buffer
                .pop_front()
                .expect("Frame is written")
                .as_frame(),
            Frame::Ping(ping::Ping),
            "Ping should be written when transmission is retransmission only"
        );
    }

    #[test]
    fn on_transmit_complete_many_transmissions_since_elicitation() {
        let mut manager =
            AckManager::new(PacketNumberSpace::ApplicationData, ack::Settings::default());
        let mut frame_buffer = OutgoingFrameBuffer::new();
        let mut write_context = MockWriteContext::new(
            time::now(),
            &mut frame_buffer,
            transmission::Constraint::None,
            transmission::Mode::Normal,
            endpoint::Type::Server,
        );
        write_context.transmission_constraint = transmission::Constraint::CongestionLimited;

        manager.ack_ranges = AckRanges::default();
        assert!(manager
            .ack_ranges
            .insert_packet_number(
                PacketNumberSpace::ApplicationData.new_packet_number(VarInt::from_u8(1)),
            )
            .is_ok());
        manager.transmission_state = AckTransmissionState::Active { retransmissions: 0 };
        manager.transmissions_since_elicitation = Counter::new(u8::max_value());

        manager.on_transmit_complete(&mut write_context);

        assert_eq!(
            manager.transmissions_since_elicitation,
            Counter::new(u8::max_value())
        );
    }

    #[test]
    #[cfg(target_pointer_width = "64")]
    fn size_of_snapshots() {
        use core::mem::size_of;
        use insta::assert_debug_snapshot;

        assert_debug_snapshot!("AckManager", size_of::<AckManager>());
    }

    #[test]
    fn client_sending_test() {
        assert_debug_snapshot!(
            "client_sending_test",
            Simulation {
                network: Network {
                    client: Application::new(
                        Endpoint::new(ack::Settings {
                            max_ack_delay: Duration::from_millis(25),
                            ack_delay_exponent: 1,
                            ..Default::default()
                        }),
                        [Duration::from_millis(5)].iter().cycle().take(100).cloned(),
                    )
                    .into(),
                    server: Application::new(
                        Endpoint::new(ack::Settings {
                            max_ack_delay: Duration::from_millis(25),
                            ack_delay_exponent: 1,
                            ..Default::default()
                        }),
                        empty(),
                    )
                    .into(),
                },
                // pass all packets unchanged
                events: empty().collect(),
                delay: Duration::from_millis(0),
            }
            .run()
        );
    }

    #[test]
    fn delayed_client_sending_test() {
        assert_debug_snapshot!(
            "delayed_client_sending_test",
            Simulation {
                network: Network {
                    client: Application::new(
                        Endpoint::new(ack::Settings {
                            max_ack_delay: Duration::from_millis(25),
                            ack_delay_exponent: 1,
                            ..Default::default()
                        }),
                        [Duration::from_millis(5)].iter().cycle().take(100).cloned(),
                    )
                    .into(),
                    server: Application::new(
                        Endpoint::new(ack::Settings {
                            max_ack_delay: Duration::from_millis(25),
                            ack_delay_exponent: 1,
                            ..Default::default()
                        }),
                        empty(),
                    )
                    .into(),
                },
                // pass all packets unchanged
                events: empty().collect(),
                // delay sending each packet by 100ms
                delay: Duration::from_millis(100),
            }
            .run()
        );
    }

    #[test]
    fn high_latency_test() {
        assert_debug_snapshot!(
            "high_latency_test",
            Simulation {
                network: Network {
                    client: Application::new(
                        Endpoint::new(ack::Settings {
                            max_ack_delay: Duration::from_millis(25),
                            ack_delay_exponent: 1,
                            ..Default::default()
                        }),
                        [Duration::from_millis(5)].iter().cycle().take(100).cloned(),
                    )
                    .into(),
                    server: Application::new(
                        Endpoint::new(ack::Settings {
                            max_ack_delay: Duration::from_millis(100),
                            ack_delay_exponent: 1,
                            ..Default::default()
                        }),
                        [Duration::from_millis(5)].iter().cycle().take(100).cloned(),
                    )
                    .into(),
                },
                // pass all packets unchanged
                events: empty().collect(),
                // delay sending each packet by 1s
                delay: Duration::from_millis(1000),
            }
            .run()
        );
    }

    #[test]
    fn lossy_network_test() {
        assert_debug_snapshot!(
            "lossy_network_test",
            Simulation {
                network: Network {
                    client: Application::new(
                        Endpoint::new(ack::Settings {
                            max_ack_delay: Duration::from_millis(25),
                            ack_delay_exponent: 1,
                            ..Default::default()
                        }),
                        [Duration::from_millis(5)].iter().cycle().take(100).cloned(),
                    )
                    .into(),
                    server: Application::new(
                        Endpoint::new(ack::Settings {
                            max_ack_delay: Duration::from_millis(100),
                            ack_delay_exponent: 1,
                            ..Default::default()
                        }),
                        [Duration::from_millis(5)].iter().cycle().take(100).cloned(),
                    )
                    .into(),
                },
                // drop every 5th packet
                events: once(NetworkEvent::Pass)
                    .cycle()
                    .take(4)
                    .chain(once(NetworkEvent::Drop))
                    .collect(),
                // delay sending each packet by 100ms
                delay: Duration::from_millis(0),
            }
            .run()
        );
    }
}


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

use core::{
    convert::TryInto,
    num::NonZeroUsize,
    ops::{Bound, Deref, DerefMut, RangeInclusive},
};
use s2n_quic_core::{
    ack::Settings,
    frame::ack,
    interval_set::{IntervalSet, RangeInclusiveIter},
    packet::number::{PacketNumber, PacketNumberRange},
    varint::VarInt,
};

#[derive(Clone, Debug)]
pub struct AckRanges(IntervalSet<PacketNumber>);

impl Default for AckRanges {
    fn default() -> Self {
        Self::new(Settings::default().ack_ranges_limit as usize)
    }
}

impl AckRanges {
    pub fn new(limit: usize) -> Self {
        let limit = NonZeroUsize::new(limit).expect("limit should be nonzero");
        Self(IntervalSet::with_limit(limit))
    }

    /// Inserts a packet number; dropping smaller values if needed
    #[inline]
    pub fn insert_packet_number_range(
        &mut self,
        pn_range: PacketNumberRange,
    ) -> Result<(), AckRangesError> {
        let interval = (
            Bound::Included(pn_range.start()),
            Bound::Included(pn_range.end()),
        );
        if self.0.insert(interval).is_ok() {
            return Ok(());
        }

        // attempt to shed the lowest packet number ranges to make room for larger ones
        match self.0.pop_min() {
            Some(min) => {
                if min < pn_range.start() {
                    let insert_res = self.0.insert(interval);
                    debug_assert!(
                        insert_res.is_ok(),
                        "min range was removed, so it should be possible to insert another range",
                    );
                    insert_res.map_err(|_| AckRangesError::RangeInsertionFailed {
                        min: pn_range.start(),
                        max: pn_range.end(),
                    })?;

                    Err(AckRangesError::LowestRangeDropped {
                        min: min.start_inclusive(),
                        max: min.end_inclusive(),
                    })
                } else {
                    // new value is smaller than min so inset it back in the front
                    let _ = self.0.insert_front(min);
                    Err(AckRangesError::RangeInsertionFailed {
                        min: pn_range.start(),
                        max: pn_range.end(),
                    })
                }
            }
            None => {
                debug_assert!(
                    false,
                    "IntervalSet should have capacity and return lowest entry"
                );
                Err(AckRangesError::RangeInsertionFailed {
                    min: pn_range.start(),
                    max: pn_range.end(),
                })
            }
        }
    }

    /// Inserts a packet number; dropping smaller values if needed
    #[inline]
    pub fn insert_packet_number(
        &mut self,
        packet_number: PacketNumber,
    ) -> Result<(), AckRangesError> {
        self.insert_packet_number_range(PacketNumberRange::new(packet_number, packet_number))
    }

    /// Returns the overall range of the ack_ranges
    #[inline]
    pub fn spread(&self) -> usize {
        match (self.min_value(), self.max_value()) {
            (Some(min), Some(max)) => {
                let min = PacketNumber::as_varint(min);
                let max = PacketNumber::as_varint(max);
                (max - min).try_into().unwrap_or(core::usize::MAX)
            }
            _ => 0,
        }
    }
}

type AckRangesIter<'a> = core::iter::Map<
    core::iter::Rev<RangeInclusiveIter<'a, PacketNumber>>,
    fn(RangeInclusive<PacketNumber>) -> RangeInclusive<VarInt>,
>;

impl<'a> ack::AckRanges for &'a AckRanges {
    type Iter = AckRangesIter<'a>;

    fn ack_ranges(&self) -> Self::Iter {
        self.0.inclusive_ranges().rev().map(|range| {
            let (start, end) = range.into_inner();
            PacketNumber::as_varint(start)..=PacketNumber::as_varint(end)
        })
    }
}

impl Deref for AckRanges {
    type Target = IntervalSet<PacketNumber>;

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

impl DerefMut for AckRanges {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum AckRangesError {
    RangeInsertionFailed {
        min: PacketNumber,
        max: PacketNumber,
    },
    LowestRangeDropped {
        min: PacketNumber,
        max: PacketNumber,
    },
}

#[cfg(test)]
mod tests {
    use bolero::check;
    use s2n_quic_core::{packet::number::PacketNumberSpace, varint};

    use super::{super::tests::packet_numbers_iter, *};

    #[test]
    fn insert_value_test() {
        let mut ack_ranges = AckRanges::new(3);
        let mut packet_numbers = packet_numbers_iter().step_by(2); // skip every other packet number

        // insert gaps up to the limit
        let pn_a = packet_numbers.next().unwrap();
        assert!(ack_ranges.insert_packet_number(pn_a).is_ok());

        let pn_b = packet_numbers.next().unwrap();
        assert!(ack_ranges.insert_packet_number(pn_b).is_ok());

        let pn_c = packet_numbers.next().unwrap();
        assert!(ack_ranges.insert_packet_number(pn_c).is_ok());

        // ensure everything was inserted
        assert_eq!(ack_ranges.interval_len(), 3);
        assert!(ack_ranges.contains(&pn_a));
        assert!(ack_ranges.contains(&pn_b));
        assert!(ack_ranges.contains(&pn_c));

        // insert a new packet number gap
        let pn_d = packet_numbers.next().unwrap();
        assert_eq!(
            ack_ranges.insert_packet_number(pn_d).err().unwrap(),
            AckRangesError::LowestRangeDropped {
                min: pn_a,
                max: pn_a
            }
        );

        // ensure the previous smaller packet number was dropped
        assert_eq!(ack_ranges.interval_len(), 3);
        assert!(!ack_ranges.contains(&pn_a));
        assert!(ack_ranges.contains(&pn_b));
        assert!(ack_ranges.contains(&pn_c));
        assert!(ack_ranges.contains(&pn_d));

        // ensure smaller values are not recorded
        assert_eq!(
            ack_ranges.insert_packet_number(pn_a).err().unwrap(),
            AckRangesError::RangeInsertionFailed {
                min: pn_a,
                max: pn_a
            }
        );
        assert_eq!(ack_ranges.interval_len(), 3);
        assert!(!ack_ranges.contains(&pn_a));
        assert!(ack_ranges.contains(&pn_b));
        assert!(ack_ranges.contains(&pn_c));
        assert!(ack_ranges.contains(&pn_d));
    }

    #[test]
    fn overlapping_range_test() {
        let mut packet_numbers = packet_numbers_iter().step_by(2); // skip every other packet number
        let mut ack_ranges = AckRanges::new(3);

        // |---a---b---c---d---e---f---g---h---i---|
        //     ^0
        //         ^-1-^
        //                 ^--2----^
        //                             ^-3-^
        //                                 ^-4-^
        //             ^--1_2--^
        //     ^--0_1--^
        let pn_a = packet_numbers.next().unwrap();
        let pn_b = packet_numbers.next().unwrap();
        let pn_c = packet_numbers.next().unwrap();
        let pn_d = packet_numbers.next().unwrap();
        let pn_e = packet_numbers.next().unwrap();
        let pn_f = packet_numbers.next().unwrap();
        let pn_g = packet_numbers.next().unwrap();
        let pn_h = packet_numbers.next().unwrap();
        let pn_i = packet_numbers.next().unwrap();
        let range_0 = PacketNumberRange::new(pn_a, pn_a);
        let range_1 = PacketNumberRange::new(pn_b, pn_c);
        let range_2 = PacketNumberRange::new(pn_d, pn_f);
        let range_3 = PacketNumberRange::new(pn_g, pn_h);
        let range_4 = PacketNumberRange::new(pn_h, pn_i);
        let range_0_1 = PacketNumberRange::new(pn_a, pn_c);
        let range_1_2 = PacketNumberRange::new(pn_c, pn_e);

        assert!(ack_ranges.insert_packet_number_range(range_1).is_ok());
        assert!(ack_ranges.insert_packet_number_range(range_2).is_ok());
        assert!(ack_ranges.insert_packet_number_range(range_3).is_ok());
        assert_eq!(ack_ranges.interval_len(), 3);

        for pn in range_1 {
            assert!(ack_ranges.contains(&pn));
        }
        for pn in range_2 {
            assert!(ack_ranges.contains(&pn));
        }
        for pn in range_3 {
            assert!(ack_ranges.contains(&pn));
        }

        // merge ranges 1 and 2
        assert!(ack_ranges.insert_packet_number_range(range_1_2).is_ok());
        assert_eq!(ack_ranges.interval_len(), 2);

        // insert range 0 at low end
        assert!(ack_ranges.insert_packet_number_range(range_0).is_ok());
        assert_eq!(ack_ranges.interval_len(), 3);

        // merge range 0_1 at low end
        assert!(ack_ranges.insert_packet_number_range(range_0_1).is_ok());
        assert_eq!(ack_ranges.interval_len(), 2);

        // merge new range at high end
        assert!(ack_ranges.insert_packet_number_range(range_4).is_ok());
        assert_eq!(ack_ranges.interval_len(), 2);
    }

    #[test]
    fn large_range_test() {
        let pn_a = PacketNumberSpace::ApplicationData.new_packet_number(VarInt::from_u32(1));
        let pn_b = PacketNumberSpace::ApplicationData
            .new_packet_number(VarInt::new(varint::MAX_VARINT_VALUE).unwrap());
        let mut ack_ranges = AckRanges::new(3);

        let range_1 = PacketNumberRange::new(pn_a, pn_b);

        assert!(ack_ranges.insert_packet_number_range(range_1).is_ok());
        assert_eq!(ack_ranges.interval_len(), 1);
    }

    #[test]
    #[cfg(target_pointer_width = "64")]
    fn size_of_snapshots() {
        use core::mem::size_of;
        use insta::assert_debug_snapshot;

        assert_debug_snapshot!("AckRanges", size_of::<AckRanges>());
    }

    #[test]
    fn insert_packet_number_range_fuzz() {
        check!()
            .with_type::<(u32, u32)>()
            .map(|(a, b)| (a.min(b), a.max(b))) // ensure valid range
            .for_each(|(a, b)| {
                let mut ack_ranges = AckRanges::new(1);

                let pn_a = PacketNumberSpace::Initial.new_packet_number(VarInt::from_u32(*a));
                let pn_b = PacketNumberSpace::Initial.new_packet_number(VarInt::from_u32(*b));
                let range_1 = PacketNumberRange::new(pn_a, pn_b);

                assert!(ack_ranges.insert_packet_number_range(range_1).is_ok());
            });
    }
}


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

use crate::{ack::ack_ranges::AckRanges, transmission};

#[derive(Clone, Copy, Default, Debug, PartialEq, Eq)]
pub enum AckTransmissionState {
    /// No ACK frames will be transmitted
    #[default]
    Disabled,

    /// An ACK frame may be transmitted, but isn't required. While in this
    /// state, the ACK Manager won't express interest in transmission, but
    /// will transmit ACK frames if other components are transmitting.
    Passive {
        /// The number of remaining tranmissions for the current ack ranges
        retransmissions: usize,
    },

    /// An ACK frame is required to be transmitted. While in this state,
    /// the ACK Manager will express interest in transmission.
    Active {
        /// The number of remaining tranmissions for the current ack ranges
        retransmissions: usize,
    },
}

impl AckTransmissionState {
    /// Returns `true` if the state is set to `Active`
    pub fn is_active(&self) -> bool {
        matches!(self, Self::Active { .. })
    }

    /// Returns `true` if ACK frames should be transmitted, either actively or passively
    pub fn should_transmit(
        &self,
        constraint: transmission::Constraint,
        mode: transmission::Mode,
        has_ranges: bool,
    ) -> bool {
        match self {
            // If we don't have any ranges to transmit, then we can't do anything
            _ if !has_ranges => false,
            // Transmitting ACK frames in probes can help the peer recover faster. Since we're sending
            // a packet here, we might as well bundle it.
            _ if !mode.is_normal() => true,
            Self::Disabled => false,
            // Only transmit acks in Passive mode if we can transmit other frames as well
            Self::Passive { .. } => constraint.can_transmit() || constraint.can_retransmit(),
            //= https://www.rfc-editor.org/rfc/rfc9002#section-7
            //# packets containing only ACK frames do not count
            //# towards bytes in flight and are not congestion controlled.
            Self::Active { .. } => true,
        }
    }

    /// Transitions the transmission to active if there are pending retransmissions
    pub fn activate(&mut self) -> &mut Self {
        if let Self::Passive { retransmissions } = *self {
            *self = AckTransmissionState::Active { retransmissions }
        }

        self
    }

    /// Notify the transmission state that pending ack ranges has updated
    pub fn on_update(&mut self, ack_ranges: &AckRanges) -> &mut Self {
        // no need to transmit anything now
        if ack_ranges.is_empty() {
            *self = AckTransmissionState::Disabled;
            return self;
        }

        //= https://www.rfc-editor.org/rfc/rfc9000#section-13.2
        //# In general, frequent feedback from a receiver improves loss and
        //# congestion response, but this has to be balanced against excessive
        //# load generated by a receiver that sends an ACK frame in response to
        //# every ack-eliciting packet.

        // By default, ACK ranges are not retransmitted
        let mut new_retransmissions = 0;

        // A high `interval_len` indicates more RX packet reordering and/or loss
        //
        // Increment retransmissions for every N intervals we are ACKing
        //
        // TODO: 2 was chosen as an initial hunch - we need to do some experiements to improve
        const INTERVAL_SCALE: usize = 2;
        new_retransmissions += ack_ranges.interval_len() / INTERVAL_SCALE;

        // A high count indicates more outstanding RX packets to be ACKed, which
        // results in higher risk of the peer retransmitting packet data which has
        // already been processed.
        //
        // Increment retransmissions for every N packets we are ACKing to increase
        // the likelihood of ACK frames being received by the peer.
        //
        // TODO: 10 was chosen as an initial hunch - we need to do some experiements to improve
        const RANGE_SCALE: usize = 10;
        new_retransmissions += ack_ranges.spread() / RANGE_SCALE;

        /// We shouldn't retransmit too many redundant ACK frames so we'll cap it
        const MAX_RETRANSMISSIONS: usize = 10;

        new_retransmissions = new_retransmissions.min(MAX_RETRANSMISSIONS);

        match self {
            Self::Active { retransmissions } => {
                *retransmissions = new_retransmissions;
            }
            Self::Passive { retransmissions } => {
                *retransmissions = new_retransmissions;
            }
            Self::Disabled => {
                *self = AckTransmissionState::Passive {
                    retransmissions: new_retransmissions,
                };
            }
        }

        self
    }

    /// Notifies the transmission state that a transmission occurred
    pub fn on_transmit(&mut self, has_ranges: bool) -> &mut Self {
        debug_assert!(
            has_ranges,
            "ACKs can only be sent when there are ranges to send"
        );

        match *self {
            Self::Active { retransmissions } | Self::Passive { retransmissions } => {
                if let Some(retransmissions) = retransmissions.checked_sub(1) {
                    *self = AckTransmissionState::Passive { retransmissions };
                } else {
                    *self = AckTransmissionState::Disabled;
                }
            }
            Self::Disabled => {
                debug_assert!(
                    has_ranges,
                    "ACKs should not be transmitted in `Disabled` state"
                );
            }
        }

        self
    }
}

impl transmission::interest::Provider for AckTransmissionState {
    #[inline]
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        if self.is_active() {
            //= https://www.rfc-editor.org/rfc/rfc9002#section-7
            //# packets containing only ACK frames do not count
            //# towards bytes in flight and are not congestion controlled
            query.on_forced()?;
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::{super::tests::packet_numbers_iter, *};
    use insta::assert_debug_snapshot;

    #[test]
    fn should_transmit_test() {
        let mut results = vec![];
        for state in &[
            AckTransmissionState::Disabled,
            AckTransmissionState::Passive { retransmissions: 1 },
            AckTransmissionState::Active { retransmissions: 1 },
        ] {
            for constraint in &[
                transmission::Constraint::None,
                transmission::Constraint::AmplificationLimited,
                transmission::Constraint::CongestionLimited,
                transmission::Constraint::RetransmissionOnly,
            ] {
                for mode in &[
                    transmission::Mode::Normal,
                    transmission::Mode::MtuProbing,
                    transmission::Mode::PathValidationOnly,
                    transmission::Mode::LossRecoveryProbing,
                ] {
                    for has_ranges in &[true, false] {
                        let res = (
                            state,
                            *constraint,
                            *mode,
                            *has_ranges,
                            state.should_transmit(*constraint, *mode, *has_ranges),
                        );
                        results.push(res);
                    }
                }
            }
        }

        assert_debug_snapshot!("should_transmit", results);
    }

    #[test]
    fn activate_test() {
        assert!(
            !AckTransmissionState::Disabled.activate().is_active(),
            "disabled state should not activate"
        );
        assert!(
            AckTransmissionState::Passive { retransmissions: 1 }
                .activate()
                .is_active(),
            "passive state should activate"
        );
        assert!(
            AckTransmissionState::Active { retransmissions: 1 }
                .activate()
                .is_active(),
            "active state should activate"
        );
    }

    #[test]
    #[should_panic]
    fn disabled_transmission_test() {
        AckTransmissionState::Disabled.on_transmit(false);
    }

    #[test]
    fn transmission_test() {
        assert_eq!(
            *AckTransmissionState::Passive { retransmissions: 0 }.on_transmit(true),
            AckTransmissionState::Disabled,
            "transmitting should transition to Disabled"
        );

        assert_eq!(
            *AckTransmissionState::Passive { retransmissions: 1 }.on_transmit(true),
            AckTransmissionState::Passive { retransmissions: 0 },
            "transmitting should decrement and stay in the same state"
        );

        assert_eq!(
            *AckTransmissionState::Disabled.on_transmit(true),
            AckTransmissionState::Disabled,
            "transmitting when disabled should continue to be disabled"
        );
    }

    #[test]
    fn update_test() {
        let mut ack_ranges = AckRanges::new(10);
        let mut packet_numbers = packet_numbers_iter().step_by(2); // skip every other packet number

        assert_eq!(
            *AckTransmissionState::Passive { retransmissions: 1 }.on_update(&ack_ranges),
            AckTransmissionState::Disabled,
            "empty ack_ranges should transition to Disabled"
        );

        assert!(ack_ranges
            .insert_packet_number(packet_numbers.next().unwrap())
            .is_ok());

        assert_eq!(
            *AckTransmissionState::Disabled.on_update(&ack_ranges),
            AckTransmissionState::Passive { retransmissions: 0 },
            "one ack range should transition to passive"
        );

        assert!(ack_ranges
            .insert_packet_number(packet_numbers.next().unwrap())
            .is_ok());
        assert!(ack_ranges
            .insert_packet_number(packet_numbers.next().unwrap())
            .is_ok());

        assert_eq!(
            *AckTransmissionState::Disabled.on_update(&ack_ranges),
            AckTransmissionState::Passive { retransmissions: 1 },
            "multiple ack ranges should transition to passive with retransmissions"
        );

        assert_eq!(
            *AckTransmissionState::Passive { retransmissions: 0 }.on_update(&ack_ranges),
            AckTransmissionState::Passive { retransmissions: 1 },
            "multiple ack ranges should update passive with retransmissions"
        );

        assert_eq!(
            *AckTransmissionState::Active { retransmissions: 0 }.on_update(&ack_ranges),
            AckTransmissionState::Active { retransmissions: 1 },
            "multiple ack ranges should update active with retransmissions"
        );
    }

    #[test]
    #[cfg(target_pointer_width = "64")]
    fn size_of_snapshots() {
        use core::mem::size_of;
        use insta::assert_debug_snapshot;

        assert_debug_snapshot!("AckTransmissionState", size_of::<AckTransmissionState>());
    }
}


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

pub use ack_manager::*;
pub use s2n_quic_core::ack::*;

mod ack_eliciting_transmission;
mod ack_manager;
pub(crate) mod ack_ranges;
mod ack_transmission_state;

#[cfg(test)]
mod tests;


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

use super::{generator::gen_duration, Endpoint, Packet};
use alloc::collections::VecDeque;
use bolero::generator::*;
use core::time::Duration;
use s2n_quic_core::{endpoint, time::Timestamp};

#[derive(Clone, Debug, TypeGenerator)]
pub struct Application {
    pub transmissions: VecDeque<ApplicationTransmission>,
    #[generator(constant(None))]
    next_transmission: Option<Timestamp>,
    pub endpoint: Endpoint,
}

impl Application {
    pub fn new<I: Iterator<Item = Duration>>(endpoint: Endpoint, transmissions: I) -> Self {
        Self {
            transmissions: transmissions.map(Into::into).collect(),
            next_transmission: None,
            endpoint,
        }
    }

    pub fn init(&mut self, now: Timestamp, endpoint_type: endpoint::Type) {
        self.init_next_transmission(now);
        self.endpoint.init(now, endpoint_type);
    }

    fn init_next_transmission(&mut self, now: Timestamp) {
        self.next_transmission = self.transmissions.pop_front().map(|t| now + t.delay);
    }

    pub fn recv(&mut self, packet: Packet) {
        self.endpoint.recv(packet)
    }

    pub fn tick(&mut self, now: Timestamp) -> Option<Packet> {
        if self.next_transmission.is_none() {
            self.init_next_transmission(now);
        }

        if self.next_transmission == Some(now) {
            self.init_next_transmission(now);
            self.endpoint.send(now)
        } else {
            self.endpoint.tick(now)
        }
    }

    pub fn timers(&self) -> impl Iterator<Item = Timestamp> {
        self.endpoint.timers().chain(self.next_transmission)
    }

    pub fn done(&mut self) {
        assert_eq!(
            self.transmissions.len(),
            0,
            "pending transmissions {:?}",
            self.transmissions
        );
        assert_eq!(self.next_transmission, None);
        self.endpoint.done();
    }
}

#[derive(Clone, Debug, Default, TypeGenerator)]
pub struct ApplicationTransmission {
    #[generator(gen_duration())]
    pub delay: Duration,
}

impl From<Duration> for ApplicationTransmission {
    fn from(delay: Duration) -> Self {
        Self { delay }
    }
}


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

use super::{generator::gen_ack_settings, Packet, TestEnvironment};
use crate::{
    ack::ack_manager::AckManager,
    contexts::WriteContext,
    path::{path_event, testing::helper_path_server},
    processed_packet::ProcessedPacket,
    transmission::interest::Provider,
};
use bolero::generator::*;
use s2n_quic_core::{
    ack, connection, endpoint, event,
    event::{testing::Publisher, IntoEvent as _},
    frame::{ack_elicitation::AckElicitation, Ack, Frame, Ping},
    inet::DatagramInfo,
    packet::number::PacketNumberSpace,
    path,
    time::{timer::Provider as _, Timestamp},
};

#[derive(Clone, Debug, TypeGenerator)]
pub struct Endpoint {
    #[generator(constant(TestEnvironment::new()))]
    pub env: TestEnvironment,

    #[generator(gen_ack_settings().map(new_ack_manager))]
    pub ack_manager: AckManager,
}

fn new_ack_manager(ack_settings: ack::Settings) -> AckManager {
    AckManager::new(PacketNumberSpace::ApplicationData, ack_settings)
}

impl Endpoint {
    pub fn new(ack_settings: ack::Settings) -> Self {
        Self {
            env: TestEnvironment::new(),
            ack_manager: new_ack_manager(ack_settings),
        }
    }

    pub fn init(&mut self, now: Timestamp, endpoint_type: endpoint::Type) {
        self.env.current_time = now;
        self.env.local_endpoint_type = endpoint_type;
    }

    pub fn recv(&mut self, packet: Packet) {
        self.env.current_time = packet.time;

        self.ack_manager.on_timeout(self.env.current_time);

        let datagram = DatagramInfo {
            ecn: packet.ecn,
            payload_len: 1200,
            timestamp: self.env.current_time,
            destination_connection_id: connection::LocalId::TEST_ID,
            source_connection_id: None,
        };

        if let Some(ack) = packet.ack {
            for ack_range in ack.ack_ranges {
                self.ack_manager
                    .on_packet_ack(datagram.timestamp, &ack_range);
            }
        }

        let packet = ProcessedPacket {
            ack_elicitation: packet.ack_elicitation,
            datagram: &datagram,
            packet_number: packet.packet_number,
            path_challenge_on_active_path: false,
            frames: 1,
            path_validation_probing: Default::default(),
            bytes_progressed: 0,
        };

        let path = helper_path_server();
        let path_id = path::Id::test_id();
        self.ack_manager.on_processed_packet(
            &packet,
            path_event!(path, path_id),
            &mut Publisher::no_snapshot(),
        );
    }

    pub fn send(&mut self, now: Timestamp) -> Option<Packet> {
        self.env.current_time = now;
        self.ack_manager.on_timeout(now);
        self.transmit(AckElicitation::Eliciting)
    }

    pub fn tick(&mut self, now: Timestamp) -> Option<Packet> {
        self.env.current_time = now;
        self.ack_manager.on_timeout(now);

        if !self.ack_manager.has_transmission_interest() {
            return None;
        }

        self.transmit(AckElicitation::NonEliciting)
    }

    pub fn timers(&self) -> impl Iterator<Item = Timestamp> {
        self.ack_manager.next_expiration().into_iter()
    }

    fn transmit(&mut self, ack_elicitation: AckElicitation) -> Option<Packet> {
        let mut context = self.env.context();
        let did_send_ack = self.ack_manager.on_transmit(&mut context);

        if ack_elicitation.is_ack_eliciting() {
            context.write_frame(&Ping);
        }

        if did_send_ack {
            self.ack_manager.on_transmit_complete(&mut context);
        }

        let ack_elicitation = context.ack_elicitation();
        let packet_number = context.packet_number();

        if self.env.sent_frames.is_empty() {
            return None;
        }

        let mut packet = Packet {
            packet_number,
            ack_elicitation,
            ecn: Default::default(),
            time: self.env.current_time,
            ack: None,
        };

        while let Some(mut frame) = self.env.sent_frames.pop_front() {
            if let Frame::Ack(ack) = frame.as_frame() {
                packet.ack = Some(Ack {
                    ack_delay: ack.ack_delay,
                    ecn_counts: ack.ecn_counts,
                    ack_ranges: ack
                        .ack_ranges()
                        .map(|ack_range| {
                            let (start, end) = ack_range.into_inner();

                            let pn_space = PacketNumberSpace::ApplicationData;
                            let start = pn_space.new_packet_number(start);
                            let end = pn_space.new_packet_number(end);

                            start..=end
                        })
                        .collect(),
                });
            }
        }

        self.env.sent_frames.flush();

        Some(packet)
    }

    pub fn done(&mut self) {
        assert!(
            !self.ack_manager.has_transmission_interest(),
            "ack manager should be in a stable state"
        );
    }
}


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

use crate::contexts::testing::{MockWriteContext, OutgoingFrameBuffer};
use s2n_quic_core::{
    endpoint,
    time::{clock::testing as time, Timestamp},
    transmission,
};

#[derive(Clone, Debug)]
pub struct TestEnvironment {
    pub sent_frames: OutgoingFrameBuffer,
    pub current_time: Timestamp,
    pub transmission_constraint: transmission::Constraint,
    pub transmission_mode: transmission::Mode,
    pub local_endpoint_type: endpoint::Type,
}

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

impl TestEnvironment {
    pub fn new() -> Self {
        let mut sent_frames = OutgoingFrameBuffer::new();

        sent_frames.set_max_packet_size(Some(1200));

        Self {
            sent_frames,
            current_time: time::now(),
            transmission_constraint: transmission::Constraint::None,
            transmission_mode: transmission::Mode::Normal,
            local_endpoint_type: endpoint::Type::Server,
        }
    }

    pub fn context(&mut self) -> MockWriteContext {
        MockWriteContext::new(
            self.current_time,
            &mut self.sent_frames,
            self.transmission_constraint,
            self.transmission_mode,
            self.local_endpoint_type,
        )
    }
}


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

use bolero::generator::*;
use core::time::Duration;
use s2n_quic_core::{ack, inet::ExplicitCongestionNotification};

pub fn gen_ack_settings() -> impl ValueGenerator<Output = ack::Settings> {
    (gen_duration(), 0..20).map(|(max_ack_delay, ack_delay_exponent)| ack::Settings {
        max_ack_delay,
        ack_delay_exponent,
        ..Default::default()
    })
}

pub fn gen_default<V: Default>() -> impl ValueGenerator<Output = V> {
    gen::<()>().map(|()| V::default())
}

pub fn gen_duration() -> impl ValueGenerator<Output = Duration> {
    ValueGenerator::map(1u16..10_000, |millis| Duration::from_millis(millis as u64))
}

pub fn gen_ecn() -> impl ValueGenerator<Output = ExplicitCongestionNotification> {
    gen::<bool>().map(|has_congestion| {
        if has_congestion {
            ExplicitCongestionNotification::Ce
        } else {
            ExplicitCongestionNotification::default()
        }
    })
}


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

use crate::ack::ack_eliciting_transmission::AckElicitingTransmission;
use s2n_quic_core::{
    packet::number::{PacketNumber, PacketNumberSpace},
    varint::VarInt,
};

/// Generates AckElicitingTransmissions from increasing packet numbers
pub fn transmissions_iter() -> impl Iterator<Item = AckElicitingTransmission> {
    packet_numbers_iter().map(|pn| AckElicitingTransmission {
        sent_in_packet: pn,
        largest_received_packet_number_acked: pn,
    })
}

/// Generates increasing packet numbers
pub fn packet_numbers_iter() -> impl Iterator<Item = PacketNumber> {
    Iterator::map(0u32.., |pn| {
        PacketNumberSpace::ApplicationData.new_packet_number(VarInt::from_u32(pn))
    })
}

pub mod application;
pub mod endpoint;
pub mod environment;
pub mod generator;
pub mod network;
pub mod network_event;
pub mod network_interface;
pub mod packet;
pub mod report;
pub mod simulation;

pub use application::*;
pub use endpoint::*;
pub use environment::*;
pub use network::*;
pub use network_event::*;
pub use network_interface::*;
pub use packet::*;
pub use report::*;
pub use simulation::*;

#[test]
fn simulation_harness_test() {
    use core::time::Duration;
    use s2n_quic_core::ack;

    let client_transmissions = 100;
    let server_transmissions = 10;

    let mut simulation = Simulation {
        network: Network {
            client: Application::new(
                Endpoint::new(ack::Settings {
                    max_ack_delay: Duration::from_millis(25),
                    ack_delay_exponent: 1,
                    ..Default::default()
                }),
                // send an ack-eliciting packet every 5ms, 100 times
                [Duration::from_millis(5)]
                    .iter()
                    .cycle()
                    .take(client_transmissions)
                    .cloned(),
            )
            .into(),
            server: Application::new(
                Endpoint::new(ack::Settings {
                    max_ack_delay: Duration::from_millis(25),
                    ack_delay_exponent: 1,
                    ..Default::default()
                }),
                // send an ack-eliciting packet every 800ms, 10 times
                [Duration::from_millis(800)]
                    .iter()
                    .cycle()
                    .take(server_transmissions)
                    .cloned(),
            )
            .into(),
        },
        // pass all packets unchanged
        events: [NetworkEvent::Pass].iter().cloned().collect(),
        // delay sending each packet by 100ms
        delay: Duration::from_millis(100),
    };

    let report = simulation.run();

    assert!(report.client.ack_eliciting_transmissions >= client_transmissions);
    assert!(report.client.dropped_transmissions == 0);

    assert!(report.server.ack_eliciting_transmissions >= server_transmissions);
    assert!(report.server.dropped_transmissions == 0);
}

/// Ack Manager Simulation
///
/// This target will generate and execute fuzz-guided simulation scenarios.
///
/// The following assertions are made:
///
/// * The program doesn't crash
///
/// * Two AckManagers talking to each other in various configurations
///   and network scenarios will always terminate (not endlessly ACK each other)
///
/// Additional checks may be implemented at some point to expand guarantees
#[test]
fn simulation_test() {
    bolero::check!()
        .with_type::<Simulation>()
        .cloned()
        .for_each(|mut simulation| {
            let _report = simulation.run();

            // TODO make assertions about report
        });
}


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

use super::{NetworkEvent, NetworkInterface, Report};
use bolero::generator::*;
use core::time::Duration;
use s2n_quic_core::{endpoint, time::Timestamp};

#[derive(Clone, Debug, TypeGenerator)]
pub struct Network {
    pub client: NetworkInterface,
    pub server: NetworkInterface,
}

impl Network {
    pub fn init(&mut self, now: Timestamp) {
        self.client.init(now, endpoint::Type::Client);
        self.server.init(now, endpoint::Type::Server);
    }

    pub fn tick<E: Iterator<Item = NetworkEvent>>(
        &mut self,
        now: Timestamp,
        events: &mut E,
        network_delay: Duration,
        report: &mut Report,
    ) {
        macro_rules! transmit {
            ($from:ident, $to:ident) => {
                if let Some(packet) = self.$from.tick(now) {
                    if let Some(mut packet) = events
                        .next()
                        .unwrap_or(NetworkEvent::Pass)
                        .process_packet(packet, &mut report.$from)
                    {
                        // add standard network delay
                        packet.time += network_delay;
                        self.$to.recv(packet);
                    }
                }
            };
        }

        transmit!(client, server);
        transmit!(server, client);
    }

    pub fn next_tick(&self) -> Option<Timestamp> {
        self.timers().min()
    }

    pub fn timers(&self) -> impl Iterator<Item = Timestamp> {
        self.client.timers().chain(self.server.timers())
    }

    pub fn done(&mut self) {
        self.client.done();
        self.server.done();
    }
}


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

use super::{
    generator::{gen_duration, gen_ecn},
    EndpointReport, Packet,
};
use bolero::generator::*;
use core::time::Duration;
use s2n_quic_core::inet::ExplicitCongestionNotification;

#[derive(Clone, Copy, Debug, TypeGenerator)]
pub enum NetworkEvent {
    Pass,
    Drop,
    Ecn(#[generator(gen_ecn())] ExplicitCongestionNotification),
    Delay(#[generator(gen_duration())] Duration),
}

impl NetworkEvent {
    pub fn process_packet(
        &self,
        mut packet: Packet,
        report: &mut EndpointReport,
    ) -> Option<Packet> {
        report.total_transmissions += 1;

        match self {
            NetworkEvent::Drop => {
                report.dropped_transmissions += 1;
                return None;
            }
            NetworkEvent::Pass => {}
            NetworkEvent::Delay(delay) => {
                packet.time += *delay;
                report.delayed_transmissions += 1;
            }
            NetworkEvent::Ecn(ecn) => {
                packet.ecn = *ecn;
                report.congested_transmissions += 1;
            }
        }

        report.processed_transmissions += 1;

        if packet.ack_elicitation.is_ack_eliciting() {
            report.ack_eliciting_transmissions += 1;
        }

        if packet.ack.is_some() {
            report.ack_transmissions += 1;
        }

        Some(packet)
    }
}


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

use super::{generator::gen_default, Application, Packet};
use alloc::collections::BTreeMap;
use bolero::generator::*;
use s2n_quic_core::{endpoint, time::Timestamp};

#[derive(Clone, Debug, TypeGenerator)]
pub struct NetworkInterface {
    pub application: Application,
    #[generator(gen_default())]
    rx_queue: BTreeMap<Timestamp, Packet>,
}

impl NetworkInterface {
    pub fn new(application: Application) -> Self {
        Self {
            application,
            rx_queue: Default::default(),
        }
    }

    pub fn init(&mut self, now: Timestamp, endpoint_type: endpoint::Type) {
        self.application.init(now, endpoint_type);
    }

    pub fn recv(&mut self, packet: Packet) {
        self.rx_queue.insert(packet.time, packet);
    }

    pub fn tick(&mut self, now: Timestamp) -> Option<Packet> {
        if let Some(packet) = self.rx_queue.remove(&now) {
            self.application.recv(packet);
        }

        self.application.tick(now)
    }

    pub fn timers(&self) -> impl Iterator<Item = Timestamp> {
        self.application
            .timers()
            .chain(self.rx_queue.keys().next().copied())
    }

    pub fn done(&mut self) {
        assert_eq!(self.rx_queue.len(), 0);
        self.application.done();
    }
}

impl From<Application> for NetworkInterface {
    fn from(value: Application) -> Self {
        NetworkInterface::new(value)
    }
}


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

use alloc::collections::VecDeque;
use core::ops::RangeInclusive;
use s2n_quic_core::{
    frame::{ack_elicitation::AckElicitation, Ack},
    inet::ExplicitCongestionNotification,
    packet::number::PacketNumber,
    time::Timestamp,
};

#[derive(Clone, Debug)]
pub struct Packet {
    pub packet_number: PacketNumber,
    pub ack_elicitation: AckElicitation,
    pub ecn: ExplicitCongestionNotification,
    pub time: Timestamp,
    pub ack: Option<Ack<VecDeque<RangeInclusive<PacketNumber>>>>,
}


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

use crate::ack::ack_ranges::AckRanges;

#[derive(Debug, Default)]
pub struct Report {
    pub client: EndpointReport,
    pub server: EndpointReport,
    pub iterations: usize,
}

#[derive(Debug, Default)]
pub struct EndpointReport {
    /// Final state of the pending AckRanges
    pub pending_ack_ranges: AckRanges,
    /// Total number of transmissions sent
    pub total_transmissions: usize,
    /// Number of transmissions that elicited an ACK
    pub ack_eliciting_transmissions: usize,
    /// Number of transmissions that contained an ACK frame
    pub ack_transmissions: usize,
    /// Number of transmissions that experienced congestion
    pub congested_transmissions: usize,
    /// Number of transmissions that were dropped by the network
    pub dropped_transmissions: usize,
    /// Number of transmissions that were delayed by the network
    pub delayed_transmissions: usize,
    /// Number of transmissions that were processed by the peer
    pub processed_transmissions: usize,
}


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

use super::{generator::gen_duration, Network, NetworkEvent, Report};
use alloc::collections::VecDeque;
use bolero::generator::*;
use s2n_quic_core::time::{clock::testing as time, Duration};

#[derive(Clone, Debug, TypeGenerator)]
pub struct Simulation {
    pub network: Network,
    pub events: VecDeque<NetworkEvent>,
    #[generator(gen_duration())]
    pub delay: Duration,
}

impl Simulation {
    pub fn run(&mut self) -> Report {
        let mut events = self.events.iter().cloned().cycle();
        let delay = self.delay;
        let mut report = Report::default();

        self.network.init(time::now());

        while let Some(now) = self.network.next_tick() {
            self.network.tick(now, &mut events, delay, &mut report);
            report.iterations += 1;
        }

        self.network.done();

        report.client.pending_ack_ranges = self
            .network
            .client
            .application
            .endpoint
            .ack_manager
            .ack_ranges
            .clone();
        report.server.pending_ack_ranges = self
            .network
            .server
            .application
            .endpoint
            .ack_manager
            .ack_ranges
            .clone();

        report
    }
}


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

//! Defines the public QUIC connection API

use crate::{
    connection::{self, ConnectionApi, OpenToken},
    stream::{ops, Stream, StreamError, StreamId},
};
use bytes::Bytes;
use core::{
    fmt,
    sync::atomic::{self, Ordering},
    task::{Context, Poll},
};
use s2n_quic_core::{
    application,
    application::ServerName,
    inet::SocketAddress,
    query::{Query, QueryMut},
    stream::StreamType,
};

/// A QUIC connection
pub struct Connection {
    /// The inner connection API implementation
    ///
    /// This uses a dynamically-dispatched interface to hide all of the connection's
    /// generic parameters and allows applications to interact with connections in a
    /// straightforward manner.
    pub(super) api: ConnectionApi,

    /// The open token associated with each connection handle.
    ///
    /// This is used to correctly track `poll_open_stream` requests.
    open_token: OpenToken,
}

impl fmt::Debug for Connection {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // TODO print interesting virtual fields
        f.debug_struct("Connection").finish()
    }
}

impl Drop for Connection {
    fn drop(&mut self) {
        debug_assert!(
            self.api.application_handle_count().load(Ordering::Acquire) > 0,
            "application_handle_count underflowed"
        );

        // Safety
        //
        // The use of Ordering and fence mirrors the `Arc` implementation in
        // the standard library.
        //
        // This fence is needed to prevent reordering of use of the data and
        // deletion of the data.  Because it is marked `Release`, the decreasing
        // of the reference count synchronizes with this `Acquire` fence. This
        // means that use of the data happens before decreasing the reference
        // count, which happens before this fence, which happens before the
        // deletion of the data.
        // https://github.com/rust-lang/rust/blob/e012a191d768adeda1ee36a99ef8b92d51920154/library/alloc/src/sync.rs#L1637

        // If the connection wasn't closed before, close it now to make sure
        // all Streams terminate.
        //
        // Only close the connection if this is the last application handle.
        // Otherwise, just drop `api`, which decrements the strong count.
        if self
            .api
            .application_handle_count()
            .fetch_sub(1, Ordering::Release)
            != 1
        {
            return;
        }

        atomic::fence(Ordering::Acquire);
        self.api.close_connection(None);
    }
}

impl Clone for Connection {
    fn clone(&self) -> Self {
        // Safety
        //
        // Using a relaxed ordering is alright here, as knowledge of the
        // original reference prevents other threads from erroneously deleting
        // the object.
        // https://github.com/rust-lang/rust/blob/e012a191d768adeda1ee36a99ef8b92d51920154/library/alloc/src/sync.rs#L1329
        self.api
            .application_handle_count()
            .fetch_add(1, Ordering::Relaxed);
        Self {
            api: self.api.clone(),
            // don't clone the open token - each instance should have its own token
            open_token: OpenToken::new(),
        }
    }
}

impl Connection {
    pub(crate) fn new(api: ConnectionApi) -> Self {
        // Safety
        //
        // Using a relaxed ordering is alright here, as knowledge of the
        // original reference prevents other threads from erroneously deleting
        // the object.
        // https://github.com/rust-lang/rust/blob/e012a191d768adeda1ee36a99ef8b92d51920154/library/alloc/src/sync.rs#L1329
        api.application_handle_count()
            .fetch_add(1, Ordering::Relaxed);
        Self {
            api,
            open_token: OpenToken::new(),
        }
    }

    /// Accepts an incoming [`Stream`]
    ///
    /// The method will return
    /// - `Poll::Ready(Ok(Some(stream, stream_type)))` if a [`Stream`] was accepted
    /// - `Poll::Ready(Ok(None))` if the connection was closed without an error
    /// - `Poll::Ready(Err(stream_error))` if no stream could be accepted due to an error
    /// - `Poll::Pending` if no new [`Stream`] of the given type was accepted by the connection yet.
    ///   In this case the caller must retry calling [`Self::poll_accept`].
    ///   For this purpose the method will save the [`core::task::Waker`]
    ///   which is provided as part of the [`Context`] parameter, and notify it
    ///   as soon as retrying the method will yield a different result.
    #[inline]
    pub fn poll_accept(
        &mut self,
        stream_type: Option<StreamType>,
        context: &Context,
    ) -> Poll<Result<Option<Stream>, connection::Error>> {
        self.api.poll_accept(&self.api, stream_type, context)
    }

    #[inline]
    pub fn poll_open_stream(
        &mut self,
        stream_type: StreamType,
        context: &Context,
    ) -> Poll<Result<Stream, connection::Error>> {
        self.api
            .poll_open_stream(&self.api, stream_type, &mut self.open_token, context)
    }

    #[inline]
    pub fn poll_request(
        &self,
        stream_id: StreamId,
        request: &mut ops::Request,
        context: Option<&Context>,
    ) -> Result<ops::Response, StreamError> {
        self.api.poll_request(stream_id, request, context)
    }

    /// Closes the Connection with the provided error code
    ///
    /// This will immediately terminate all outstanding streams.
    #[inline]
    pub fn close(&self, error_code: application::Error) {
        self.api.close_connection(Some(error_code));
    }

    #[inline]
    pub fn server_name(&self) -> Result<Option<ServerName>, connection::Error> {
        self.api.server_name()
    }

    #[inline]
    pub fn application_protocol(&self) -> Result<Bytes, connection::Error> {
        self.api.application_protocol()
    }

    #[inline]
    pub fn id(&self) -> u64 {
        self.api.id()
    }

    #[inline]
    pub fn ping(&self) -> Result<(), connection::Error> {
        self.api.ping()
    }

    pub fn keep_alive(&self, enabled: bool) -> Result<(), connection::Error> {
        self.api.keep_alive(enabled)
    }

    #[inline]
    pub fn local_address(&self) -> Result<SocketAddress, connection::Error> {
        self.api.local_address()
    }

    #[inline]
    pub fn remote_address(&self) -> Result<SocketAddress, connection::Error> {
        self.api.remote_address()
    }

    #[inline]
    pub fn query_event_context(&self, query: &mut dyn Query) -> Result<(), connection::Error> {
        self.api.query_event_context(query)
    }

    #[inline]
    pub fn query_event_context_mut(
        &self,
        query: &mut dyn QueryMut,
    ) -> Result<(), connection::Error> {
        self.api.query_event_context_mut(query)
    }

    #[inline]
    pub fn datagram_mut(&self, query: &mut dyn QueryMut) -> Result<(), connection::Error> {
        self.api.datagram_mut(query)
    }
}


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

//! Defines the trait for types which provide the public Connection and Stream API
//! via dynamic dispatch

use crate::{
    connection,
    stream::{Stream, StreamError},
};
use alloc::sync::Arc;
use bytes::Bytes;
use core::{
    sync::atomic::AtomicUsize,
    task::{Context, Poll},
};
use s2n_quic_core::{
    application,
    application::ServerName,
    inet::SocketAddress,
    query::{Query, QueryMut},
    stream::{ops, StreamId, StreamType},
};

/// A dynamically dispatched connection API
pub(crate) type ConnectionApi = Arc<dyn ConnectionApiProvider>;

/// The trait for types which provide the public Connection and Stream API via
/// dynamic dispatch
pub(crate) trait ConnectionApiProvider: Sync + Send {
    fn application_handle_count(&self) -> &AtomicUsize;

    fn poll_request(
        &self,
        stream_id: StreamId,
        request: &mut ops::Request,
        context: Option<&Context>,
    ) -> Result<ops::Response, StreamError>;

    fn poll_accept(
        &self,
        arc_self: &Arc<dyn ConnectionApiProvider>,
        stream_type: Option<StreamType>,
        context: &Context,
    ) -> Poll<Result<Option<Stream>, connection::Error>>;

    fn poll_open_stream(
        &self,
        arc_self: &Arc<dyn ConnectionApiProvider>,
        stream_type: StreamType,
        open_token: &mut connection::OpenToken,
        context: &Context,
    ) -> Poll<Result<Stream, connection::Error>>;

    fn close_connection(&self, code: Option<application::Error>);

    fn server_name(&self) -> Result<Option<ServerName>, connection::Error>;

    fn application_protocol(&self) -> Result<Bytes, connection::Error>;

    fn id(&self) -> u64;

    fn ping(&self) -> Result<(), connection::Error>;

    fn keep_alive(&self, enabled: bool) -> Result<(), connection::Error>;

    fn local_address(&self) -> Result<SocketAddress, connection::Error>;

    fn remote_address(&self) -> Result<SocketAddress, connection::Error>;

    fn query_event_context(&self, query: &mut dyn Query) -> Result<(), connection::Error>;

    fn query_event_context_mut(&self, query: &mut dyn QueryMut) -> Result<(), connection::Error>;

    fn datagram_mut(&self, query: &mut dyn QueryMut) -> Result<(), connection::Error>;
}


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

use crate::{
    connection::finalization,
    endpoint,
    path::Path,
    transmission::{self, interest::Provider as _},
};
use bytes::Bytes;
use core::{task::Poll, time::Duration};
use s2n_quic_core::{
    counter::{self, Counter},
    event,
    inet::ExplicitCongestionNotification,
    io::tx,
    time::{timer, Timer, Timestamp},
};

#[derive(Debug, Default)]
pub struct CloseSender {
    state: State,
}

impl CloseSender {
    pub fn close(&mut self, packet: Bytes, timeout: Duration, now: Timestamp) {
        debug_assert!(matches!(self.state, State::Idle));

        let mut close_timer = Timer::default();
        close_timer.set(now + timeout);

        self.state = State::Closing {
            packet,
            transmission: TransmissionState::Transmitting,
            close_timer,
            limiter: Limiter::default(),
        };
    }

    pub fn on_timeout(&mut self, now: Timestamp) -> Poll<()> {
        self.state.on_timeout(now)
    }

    pub fn on_datagram_received(&mut self, rtt: Duration, now: Timestamp) {
        self.state.on_datagram_received(rtt, now);
    }

    pub fn transmission<'a, Config: endpoint::Config, Pub: event::ConnectionPublisher>(
        &'a mut self,
        path: &'a mut Path<Config>,
        timestamp: Timestamp,
        publisher: &'a mut Pub,
    ) -> Transmission<'a, Config, Pub> {
        debug_assert!(
            self.has_transmission_interest(),
            "transmission should only be called when transmission interest is expressed"
        );

        if let State::Closing {
            packet,
            transmission,
            ..
        } = &mut self.state
        {
            Transmission {
                packet,
                transmission,
                path,
                timestamp,
                publisher,
            }
        } else {
            unreachable!(
                "transmission should only be called when close sender has transmission interest"
            )
        }
    }
}

impl timer::Provider for CloseSender {
    #[inline]
    fn timers<Q: timer::Query>(&self, query: &mut Q) -> timer::Result {
        self.state.timers(query)?;

        Ok(())
    }
}

impl finalization::Provider for CloseSender {
    fn finalization_status(&self) -> finalization::Status {
        match &self.state {
            State::Idle => finalization::Status::Idle,
            State::Closing { .. } => finalization::Status::Draining,
            State::Closed => finalization::Status::Final,
        }
    }
}

impl transmission::interest::Provider for CloseSender {
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        if matches!(
            self.state,
            State::Closing {
                transmission: TransmissionState::Transmitting,
                ..
            }
        ) {
            //= https://www.rfc-editor.org/rfc/rfc9002#section-3
            //# Packets containing frames besides ACK or CONNECTION_CLOSE frames
            //# count toward congestion control limits and are considered to be in
            //# flight.

            // this packet only contains a CONNECTION_CLOSE so bypass the CC
            query.on_forced()?;
        }

        Ok(())
    }
}

pub struct Transmission<'a, Config: endpoint::Config, Pub: event::ConnectionPublisher> {
    packet: &'a Bytes,
    transmission: &'a mut TransmissionState,
    path: &'a mut Path<Config>,
    timestamp: Timestamp,
    publisher: &'a mut Pub,
}

impl<'a, Config: endpoint::Config, Pub: event::ConnectionPublisher> tx::Message
    for Transmission<'a, Config, Pub>
{
    type Handle = Config::PathHandle;

    #[inline]
    fn path_handle(&self) -> &Self::Handle {
        &self.path.handle
    }

    #[inline]
    fn ecn(&mut self) -> ExplicitCongestionNotification {
        self.path
            .ecn_controller
            .ecn(transmission::Mode::Normal, self.timestamp)
    }

    #[inline]
    fn ipv6_flow_label(&mut self) -> u32 {
        0
    }

    #[inline]
    fn can_gso(&self, segment_len: usize, _segment_count: usize) -> bool {
        segment_len >= self.packet.len()
    }

    #[inline]
    fn delay(&mut self) -> Duration {
        Duration::ZERO
    }

    #[inline]
    fn write_payload(
        &mut self,
        mut buffer: tx::PayloadBuffer,
        gso_offset: usize,
    ) -> Result<usize, tx::Error> {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-10.2.1
        //# |  Note: Allowing retransmission of a closing packet is an
        //# |  exception to the requirement that a new packet number be used
        //# |  for each packet; see Section 12.3.  Sending new packet numbers
        //# |  is primarily of advantage to loss recovery and congestion
        //# |  control, which are not expected to be relevant for a closed
        //# |  connection.  Retransmitting the final packet requires less
        //# |  state.
        let len = buffer.write(self.packet)?;

        self.path.on_bytes_transmitted(len);
        *self.transmission = TransmissionState::Idle;

        self.publisher
            .on_datagram_sent(event::builder::DatagramSent {
                len: len as u16,
                gso_offset,
            });

        Ok(len)
    }
}

#[derive(Debug)]
enum State {
    Idle,
    Closing {
        packet: Bytes,
        limiter: Limiter,
        transmission: TransmissionState,
        close_timer: Timer,
    },
    Closed,
}

impl Default for State {
    fn default() -> Self {
        Self::Idle
    }
}

impl State {
    pub fn on_timeout(&mut self, now: Timestamp) -> Poll<()> {
        match self {
            Self::Idle => Poll::Pending,
            Self::Closing {
                close_timer,
                limiter,
                transmission,
                ..
            } => {
                if close_timer.poll_expiration(now).is_ready() {
                    *self = Self::Closed;
                    return Poll::Ready(());
                }

                if limiter.on_timeout(now).is_ready() {
                    *transmission = TransmissionState::Transmitting;
                }

                Poll::Pending
            }
            Self::Closed => Poll::Ready(()),
        }
    }

    pub fn on_datagram_received(&mut self, rtt: Duration, now: Timestamp) {
        if let Self::Closing { limiter, .. } = self {
            limiter.on_datagram_received(rtt, now);
        }
    }
}

impl timer::Provider for State {
    #[inline]
    fn timers<Q: timer::Query>(&self, query: &mut Q) -> timer::Result {
        if let Self::Closing {
            close_timer,
            limiter,
            ..
        } = self
        {
            close_timer.timers(query)?;
            limiter.timers(query)?;
        };

        Ok(())
    }
}

#[derive(Debug)]
enum TransmissionState {
    Idle,
    Transmitting,
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-10.2.1
//# An endpoint SHOULD limit the rate at which it generates packets in
//# the closing state.  For instance, an endpoint could wait for a
//# progressively increasing number of received packets or amount of time
//# before responding to received packets.
#[derive(Debug)]
struct Limiter {
    factor: Counter<u8, counter::Saturating>,
    received: Counter<u8, counter::Saturating>,
    debounce: Timer,
}

impl Default for Limiter {
    fn default() -> Self {
        Self {
            factor: Counter::new(1),
            received: Counter::new(0),
            debounce: Timer::default(),
        }
    }
}

impl Limiter {
    pub fn on_timeout(&mut self, now: Timestamp) -> Poll<()> {
        self.debounce.poll_expiration(now)
    }

    pub fn on_datagram_received(&mut self, rtt: Duration, now: Timestamp) {
        if self.debounce.is_armed() {
            return;
        }

        self.received += 1;

        if self.received >= self.factor {
            self.received = Counter::new(0);
            self.factor += self.factor;
            self.debounce.set(now + rtt);
        }
    }
}

impl timer::Provider for Limiter {
    #[inline]
    fn timers<Q: timer::Query>(&self, query: &mut Q) -> timer::Result {
        self.debounce.timers(query)?;

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::path::testing::helper_path_server;
    use s2n_quic_core::{
        event::testing::Publisher,
        io::tx::Message as _,
        path::MINIMUM_MTU,
        time::{clock::testing as time, testing::Clock, timer::Provider as _, Clock as _},
    };

    static PACKET: Bytes = Bytes::from_static(b"CLOSE");

    #[test]
    fn model_test() {
        use bolero::generator::*;

        let close_time = (100..=3000).map_gen(Duration::from_millis);
        let rtt = close_time.clone().map_gen(|t| t / 3);
        let packet_durations = (0..=5000).map_gen(Duration::from_millis);
        let packet_sizes = 1..=9000;
        let events = gen::<Vec<_>>()
            .with()
            .values((packet_durations, packet_sizes));

        bolero::check!()
            .with_generator((close_time, rtt, events, gen()))
            .for_each(|(close_time, rtt, events, is_validated)| {
                let mut sender = CloseSender::default();
                let mut clock = Clock::default();
                let mut path = helper_path_server();
                let mut buffer = [0; MINIMUM_MTU as usize];
                let mut transmission_count = 0usize;
                let mut publisher = Publisher::no_snapshot();

                if *is_validated {
                    // simulate receiving a handshake packet to force path validation
                    path.on_handshake_packet();
                } else {
                    // give the path some initial credits
                    path.on_bytes_received(MINIMUM_MTU as usize);
                }

                path.on_closing();
                sender.close(PACKET.clone(), *close_time, clock.get_time());

                // transmit an initial packet
                assert!(sender.can_transmit(path.transmission_constraint()));
                let now = time::now();
                let _ = sender
                    .transmission(&mut path, now, &mut publisher)
                    .write_payload(tx::PayloadBuffer::new(&mut buffer), 0);

                for (gap, packet_size) in events {
                    // get the next timer event
                    let mut gap = *gap;
                    if let Some(expiration) = sender.next_expiration() {
                        gap = gap.min(expiration - clock.get_time());
                    }
                    clock.inc_by(gap);

                    // notify that we've received an incoming packet
                    sender.on_datagram_received(*rtt, clock.get_time());
                    path.on_bytes_received(*packet_size);

                    // try to send multiple times to ensure we only
                    // send a single packet
                    let transmission_count_before = transmission_count;
                    for _ in 0..3 {
                        let interest = sender.get_transmission_interest();
                        if interest.can_transmit(path.transmission_constraint()) {
                            let _ = sender
                                .transmission(&mut path, now, &mut publisher)
                                .write_payload(tx::PayloadBuffer::new(&mut buffer), 0);
                            transmission_count += 1;
                        }
                    }
                    assert!(transmission_count - transmission_count_before <= 1);
                }

                // make sure we eventually clean up the sender
                clock.inc_by(*close_time);
                assert!(sender.on_timeout(clock.get_time()).is_ready());

                if events.is_empty() {
                    assert_eq!(transmission_count, 0);
                } else {
                    assert!(
                        transmission_count < events.len(),
                        "transmission count should never exceed the number of events"
                    );
                }
            })
    }

    #[test]
    fn limiter_test() {
        let mut limiter = Limiter::default();
        let mut clock = Clock::default();
        let rtt = Duration::from_millis(250);

        for count in (0..10).map(|v| 2usize.pow(v)) {
            for _ in 0..(count - 1) {
                limiter.on_datagram_received(rtt, clock.get_time());
            }

            // if count doesn't saturate the counter, make sure we're not armed
            if count < 256 {
                assert!(limiter.on_timeout(clock.get_time()).is_pending());
                limiter.on_datagram_received(rtt, clock.get_time());
            }

            assert!(limiter.on_timeout(clock.get_time()).is_pending());
            clock.inc_by(rtt);
            assert!(limiter.on_timeout(clock.get_time()).is_ready());
        }
    }
}