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

//! `StreamContainer` is a container for all Streams. It manages the permanent
//! map of all active Streams, as well as a variety of dynamic Stream lists.

// hide warnings from the intrusive_collections crate
#![allow(unknown_lints, clippy::non_send_fields_in_send_ty)]

use crate::{
    stream,
    stream::{stream_impl::StreamTrait, stream_interests::StreamInterests},
    transmission,
};
use alloc::rc::Rc;
use core::{cell::RefCell, ops::Deref};
use intrusive_collections::{
    intrusive_adapter, KeyAdapter, LinkedList, LinkedListLink, RBTree, RBTreeLink,
};
use s2n_quic_core::{stream::StreamId, time::timer};

// Intrusive list adapter for managing the list of `done` streams
intrusive_adapter!(DoneStreamsAdapter<S> = Rc<StreamNode<S>>: StreamNode<S> {
    done_streams_link: LinkedListLink
});

// Intrusive list adapter for managing the list of `waiting_for_frame_delivery` streams
intrusive_adapter!(WaitingForFrameDeliveryAdapter<S> = Rc<StreamNode<S>>: StreamNode<S> {
    waiting_for_frame_delivery_link: LinkedListLink
});

// Intrusive list adapter for managing the list of
// `waiting_for_transmission` streams
intrusive_adapter!(WaitingForTransmissionAdapter<S> = Rc<StreamNode<S>>: StreamNode<S> {
    waiting_for_transmission_link: LinkedListLink
});

// Intrusive list adapter for managing the list of
// `waiting_for_retransmission` streams
intrusive_adapter!(WaitingForRetransmissionAdapter<S> = Rc<StreamNode<S>>: StreamNode<S> {
    waiting_for_retransmission_link: LinkedListLink
});

// Intrusive list adapter for managing the list of
// `waiting_for_connection_flow_control_credits` streams
intrusive_adapter!(WaitingForConnectionFlowControlCreditsAdapter<S> = Rc<StreamNode<S>>: StreamNode<S> {
    waiting_for_connection_flow_control_credits_link: LinkedListLink
});

// Intrusive list adapter for managing the list of
// `waiting_for_stream_flow_control_credits` streams
intrusive_adapter!(WaitingForStreamFlowControlCreditsAdapter<S> = Rc<StreamNode<S>>: StreamNode<S> {
    waiting_for_stream_flow_control_credits_link: LinkedListLink
});

// Intrusive red black tree adapter for managing all streams in a tree for
// lookup by Stream ID
intrusive_adapter!(StreamTreeAdapter<S> = Rc<StreamNode<S>>: StreamNode<S> {
    tree_link: RBTreeLink
});

/// A wrapper around a `Stream` implementation which allows to insert the
/// it in multiple intrusive collections. The collections into which the `Stream`
/// gets inserted are referenced inside this `StreamNode`.
struct StreamNode<S> {
    /// This contains the actual implementation of the `Stream`
    inner: RefCell<S>,
    /// Allows the Stream to be part of the `stream_map` collection
    tree_link: RBTreeLink,
    /// Allows the Stream to be part of the `done_streams` collection
    done_streams_link: LinkedListLink,
    /// Allows the Stream to be part of the `waiting_for_frame_delivery` collection
    waiting_for_frame_delivery_link: LinkedListLink,
    /// Allows the Stream to be part of the `waiting_for_transmission` collection
    waiting_for_transmission_link: LinkedListLink,
    /// Allows the Stream to be part of the `waiting_for_transmission` collection
    waiting_for_retransmission_link: LinkedListLink,
    /// Allows the Stream to be part of the `waiting_for_connection_flow_control_credits` collection
    waiting_for_connection_flow_control_credits_link: LinkedListLink,
    /// Allows the Stream to be part of the `waiting_for_stream_flow_control_credits` collection
    waiting_for_stream_flow_control_credits_link: LinkedListLink,
}

impl<S> StreamNode<S> {
    /// Creates a new `StreamNode` which wraps the given Stream implementation of type `S`
    pub fn new(stream_impl: S) -> StreamNode<S> {
        StreamNode {
            inner: RefCell::new(stream_impl),
            tree_link: RBTreeLink::new(),
            done_streams_link: LinkedListLink::new(),
            waiting_for_frame_delivery_link: LinkedListLink::new(),
            waiting_for_transmission_link: LinkedListLink::new(),
            waiting_for_retransmission_link: LinkedListLink::new(),
            waiting_for_connection_flow_control_credits_link: LinkedListLink::new(),
            waiting_for_stream_flow_control_credits_link: LinkedListLink::new(),
        }
    }
}

// This is required to build an intrusive `RBTree` of `StreamNode`s which
// utilizes `StreamId`s as a key.
impl<'a, S: StreamTrait> KeyAdapter<'a> for StreamTreeAdapter<S> {
    type Key = StreamId;

    fn get_key(&self, x: &'a StreamNode<S>) -> StreamId {
        x.inner.borrow().stream_id()
    }
}

/// Obtains a `Rc<StreamNode>` from a `&StreamNode`.
///
/// This method is only safe to be called if the `StreamNode` is known to be
/// stored inside a `Rc`.
unsafe fn stream_node_rc_from_ref<S>(stream_node: &StreamNode<S>) -> Rc<StreamNode<S>> {
    // In order to be able to to get a `Rc` we construct a temporary `Rc`
    // from it using the `Rc::from_raw` API and clone the `Rc`.
    // The temporary `Rc` must be released without calling `drop`,
    // because this would decrement and thereby invalidate the refcount
    // (which wasn't changed by calling `Rc::from_raw`).
    let temp_node_ptr: core::mem::ManuallyDrop<Rc<StreamNode<S>>> = core::mem::ManuallyDrop::new(
        Rc::<StreamNode<S>>::from_raw(stream_node as *const StreamNode<S>),
    );
    temp_node_ptr.deref().clone()
}

/// Contains all secondary lists of Streams.
///
/// A Stream can be a member in any of those, in addition to being a member of
/// `StreamContainer::stream_map`.
struct InterestLists<S> {
    /// Streams which have been finalized
    done_streams: LinkedList<DoneStreamsAdapter<S>>,
    /// Streams which are waiting for packet acknowledgements and
    /// packet loss notifications
    waiting_for_frame_delivery: LinkedList<WaitingForFrameDeliveryAdapter<S>>,
    /// Streams which need to transmit data
    waiting_for_transmission: LinkedList<WaitingForTransmissionAdapter<S>>,
    /// Streams which need to transmit data
    waiting_for_retransmission: LinkedList<WaitingForRetransmissionAdapter<S>>,
    /// Streams which are blocked on transmission due to waiting on the
    /// connection flow control window to increase
    waiting_for_connection_flow_control_credits:
        LinkedList<WaitingForConnectionFlowControlCreditsAdapter<S>>,
    /// Streams which are blocked on transmission due to waiting on the
    /// stream flow control window to increase
    waiting_for_stream_flow_control_credits:
        LinkedList<WaitingForStreamFlowControlCreditsAdapter<S>>,
}

impl<S: StreamTrait> InterestLists<S> {
    fn new() -> Self {
        Self {
            done_streams: LinkedList::new(DoneStreamsAdapter::new()),
            waiting_for_frame_delivery: LinkedList::new(WaitingForFrameDeliveryAdapter::new()),
            waiting_for_transmission: LinkedList::new(WaitingForTransmissionAdapter::new()),
            waiting_for_retransmission: LinkedList::new(WaitingForRetransmissionAdapter::new()),
            waiting_for_connection_flow_control_credits: LinkedList::new(
                WaitingForConnectionFlowControlCreditsAdapter::new(),
            ),
            waiting_for_stream_flow_control_credits: LinkedList::new(
                WaitingForStreamFlowControlCreditsAdapter::new(),
            ),
        }
    }

    /// Update all interest lists based on latest interest reported by a Node
    fn update_interests(
        &mut self,
        node: &Rc<StreamNode<S>>,
        interests: StreamInterests,
        result: StreamContainerIterationResult,
    ) -> bool {
        // Note that all comparisons start by checking whether the stream is
        // already part of the given list. This is required in order for the
        // following operation to be safe. Inserting an element in a list while
        // it is already part of a (different) list can panic. Trying to remove
        // an element from a list while it is not actually part of the list
        // is undefined.

        macro_rules! sync_interests {
            ($interest:expr, $link_name:ident, $list_name:ident) => {
                if $interest != node.$link_name.is_linked() {
                    if $interest {
                        if matches!(result, StreamContainerIterationResult::Continue) {
                            self.$list_name.push_back(node.clone());
                        } else {
                            self.$list_name.push_front(node.clone());
                        }
                    } else {
                        // Safety: We know that the node is only ever part of this list.
                        // While elements are in temporary lists, they always get unlinked
                        // from those temporary lists while their interest is updated.
                        let mut cursor = unsafe {
                            self.$list_name
                                .cursor_mut_from_ptr(node.deref() as *const StreamNode<S>)
                        };
                        cursor.remove();
                    }
                }
                debug_assert_eq!($interest, node.$link_name.is_linked());
            };
        }

        sync_interests!(
            interests.delivery_notifications,
            waiting_for_frame_delivery_link,
            waiting_for_frame_delivery
        );
        sync_interests!(
            matches!(interests.transmission, transmission::Interest::NewData),
            waiting_for_transmission_link,
            waiting_for_transmission
        );
        sync_interests!(
            matches!(interests.transmission, transmission::Interest::LostData),
            waiting_for_retransmission_link,
            waiting_for_retransmission
        );
        sync_interests!(
            interests.connection_flow_control_credits,
            waiting_for_connection_flow_control_credits_link,
            waiting_for_connection_flow_control_credits
        );
        sync_interests!(
            interests.stream_flow_control_credits,
            waiting_for_stream_flow_control_credits_link,
            waiting_for_stream_flow_control_credits
        );

        if !interests.retained != node.done_streams_link.is_linked() {
            if !interests.retained {
                self.done_streams.push_back(node.clone());
            } else {
                panic!("Done streams should never report not done later");
            }
            true
        } else {
            false
        }
    }
}

/// A collection of all intrusive lists Streams are part of.
///
/// The container will automatically update the membership of a `Stream` in a
/// variety of interest lists after each interaction with the `Stream`.
///
/// The Stream container can be interacted with in 2 fashions:
/// - The `with_stream()` method allows users to obtain a mutable reference to
///   a single `Stream`. After the interaction was completed, the `Stream` will
///   be queried for its interests again.
/// - There exist a variety of iteration methods, which allow to iterate over
///   all or a subset of streams in each interest list.
pub struct StreamContainer<S> {
    /// Streams organized as a tree, for lookup by Stream ID
    stream_map: RBTree<StreamTreeAdapter<S>>,
    /// The number of streams which are tracked by the Container.
    /// This needs to be in-sync with Streams that get inserted into `stream_map`.
    nr_active_streams: usize,
    /// Additional interest lists in which Streams will be placed dynamically
    interest_lists: InterestLists<S>,
}

impl<S> core::fmt::Debug for StreamContainer<S> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> Result<(), core::fmt::Error> {
        f.debug_struct("StreamContainer")
            .field("nr_active_streams", &self.nr_active_streams)
            .finish()
    }
}

macro_rules! iterate_uninterruptible {
    ($sel:ident, $list_name:tt, $link_name:ident, $controller:ident, $func:ident) => {
        for stream in $sel.interest_lists.$list_name.take() {
            debug_assert!(!stream.$link_name.is_linked());

            let interests = {
                let mut mut_stream = stream.inner.borrow_mut();
                $func(&mut *mut_stream);
                mut_stream.get_stream_interests()
            };

            $sel.interest_lists.update_interests(
                &stream,
                interests,
                StreamContainerIterationResult::Continue,
            );
        }

        if !$sel.interest_lists.done_streams.is_empty() {
            $sel.finalize_done_streams($controller);
        }
    };
}

macro_rules! iterate_interruptible {
    ($sel:ident, $list_name:tt, $link_name:ident, $controller:ident, $func:ident) => {
        let mut extracted_list = $sel.interest_lists.$list_name.take();
        let mut cursor = extracted_list.front_mut();

        while let Some(stream) = cursor.remove() {
            // Note that while we iterate over the intrusive lists here
            // `stream` is part of no list anymore, since it also got dropped
            // from list that is described by the `cursor`.
            debug_assert!(!stream.$link_name.is_linked());
            let mut mut_stream = stream.inner.borrow_mut();
            let result = $func(&mut *mut_stream);

            // Update the interests after the interaction
            let interests = mut_stream.get_stream_interests();
            $sel.interest_lists
                .update_interests(&stream, interests, result);

            match result {
                StreamContainerIterationResult::BreakAndInsertAtBack => {
                    $sel.interest_lists
                        .$list_name
                        .front_mut()
                        .splice_after(extracted_list);
                    break;
                }
                StreamContainerIterationResult::Continue => {}
            }
        }

        if !$sel.interest_lists.done_streams.is_empty() {
            $sel.finalize_done_streams($controller);
        }
    };
}

impl<S: StreamTrait> StreamContainer<S> {
    /// Creates a new `StreamContainer`
    pub fn new() -> Self {
        Self {
            stream_map: RBTree::new(StreamTreeAdapter::new()),
            nr_active_streams: 0,
            interest_lists: InterestLists::new(),
        }
    }

    /// Insert a new Stream into the container
    pub fn insert_stream(&mut self, stream: S) {
        // Even though it likely might have none, it seems like it
        // would be better to avoid future bugs
        let interests = stream.get_stream_interests();

        let new_stream = Rc::new(StreamNode::new(stream));

        self.interest_lists.update_interests(
            &new_stream,
            interests,
            StreamContainerIterationResult::Continue,
        );

        self.stream_map.insert(new_stream);
        self.nr_active_streams += 1;
    }

    /// Returns the amount of streams which are tracked by the `StreamContainer`
    pub fn nr_active_streams(&self) -> usize {
        self.nr_active_streams
    }

    /// Returns true if the container contains a Stream with the given ID
    pub fn contains(&self, stream_id: StreamId) -> bool {
        !self.stream_map.find(&stream_id).is_null()
    }

    /// Looks up the `Stream` with the given ID and executes the provided function
    /// on it.
    ///
    /// After the transaction with the `Stream` had been completed, the `Stream`
    /// will get queried for its new interests, and all lists will be updated
    /// according to those.
    ///
    /// The `stream::Controller` will be notified of streams that have been
    /// closed to allow for further streams to be opened.
    ///
    /// `Stream`s which signal finalization interest will be removed from the
    /// `StreamContainer`.
    pub fn with_stream<F, R>(
        &mut self,
        stream_id: StreamId,
        controller: &mut stream::Controller,
        func: F,
    ) -> Option<R>
    where
        F: FnOnce(&mut S) -> R,
    {
        let node_ptr: Rc<StreamNode<S>>;
        let result: R;
        let interests;

        // This block is required since we mutably borrow `self` inside the
        // block in order to obtain a Stream reference and to executing the
        // provided method.
        // We need to release the borrow in order to be able to update the
        // Streams interests after having executed the method.
        {
            let node = self.stream_map.find(&stream_id).get()?;

            // We have to obtain an `Rc<StreamNode>` in order to be able to
            // perform interest updates later on. However the intrusive tree
            // API only provides us a raw reference.
            // Safety: We know that all of our StreamNode's are stored in
            // `Rc` pointers.
            node_ptr = unsafe { stream_node_rc_from_ref(node) };

            let stream: &mut S = &mut node.inner.borrow_mut();
            result = func(stream);
            interests = stream.get_stream_interests();
        }

        // Update the interest lists after the interactions and then remove
        // all finalized streams
        if self.interest_lists.update_interests(
            &node_ptr,
            interests,
            StreamContainerIterationResult::Continue,
        ) {
            self.finalize_done_streams(controller);
        }

        Some(result)
    }

    /// Removes all Streams in the `done` state from the `StreamManager`.
    ///
    /// The `stream::Controller` will be notified of streams that have been
    /// closed to allow for further streams to be opened.
    pub fn finalize_done_streams(&mut self, controller: &mut stream::Controller) {
        for stream in self.interest_lists.done_streams.take() {
            // Remove the Stream from `stream_map`
            let mut cursor = self.stream_map.find_mut(&stream.inner.borrow().stream_id());
            let remove_result = cursor.remove();
            debug_assert!(remove_result.is_some());
            self.nr_active_streams -= 1;

            // And remove the Stream from all other interest lists it might be
            // part of.
            let stream_ptr = &*stream as *const StreamNode<S>;

            macro_rules! remove_stream_from_list {
                ($list_name:ident, $link_name:ident) => {
                    if stream.$link_name.is_linked() {
                        // Safety: We know that the Stream is part of the list,
                        // because it is linked, and we never place Streams in
                        // other lists when `finalize_done_streams` is called.
                        let mut cursor = unsafe {
                            self.interest_lists
                                .$list_name
                                .cursor_mut_from_ptr(stream_ptr)
                        };
                        let remove_result = cursor.remove();
                        debug_assert!(remove_result.is_some());
                    }
                };
            }

            remove_stream_from_list!(waiting_for_frame_delivery, waiting_for_frame_delivery_link);
            remove_stream_from_list!(waiting_for_transmission, waiting_for_transmission_link);
            remove_stream_from_list!(waiting_for_retransmission, waiting_for_retransmission_link);
            remove_stream_from_list!(
                waiting_for_connection_flow_control_credits,
                waiting_for_connection_flow_control_credits_link
            );
            remove_stream_from_list!(
                waiting_for_stream_flow_control_credits,
                waiting_for_stream_flow_control_credits_link
            );

            controller.on_close_stream(stream.inner.borrow().stream_id());
        }
    }

    /// Iterates over all `Stream`s which are waiting for frame delivery,
    /// and executes the given function on each `Stream`
    ///
    /// The `stream::Controller` will be notified of streams that have been
    /// closed to allow for further streams to be opened.
    pub fn iterate_frame_delivery_list<F>(
        &mut self,
        controller: &mut stream::Controller,
        mut func: F,
    ) where
        F: FnMut(&mut S),
    {
        iterate_uninterruptible!(
            self,
            waiting_for_frame_delivery,
            waiting_for_frame_delivery_link,
            controller,
            func
        );
    }

    /// Iterates over all `Stream`s which waiting for connection flow control
    /// credits, and executes the given function on each `Stream`
    ///
    /// The `stream::Controller` will be notified of streams that have been
    /// closed to allow for further streams to be opened.
    pub fn iterate_connection_flow_credits_list<F>(
        &mut self,
        controller: &mut stream::Controller,
        mut func: F,
    ) where
        F: FnMut(&mut S) -> StreamContainerIterationResult,
    {
        iterate_interruptible!(
            self,
            waiting_for_connection_flow_control_credits,
            waiting_for_connection_flow_control_credits_link,
            controller,
            func
        );
    }

    /// Iterates over all `Stream`s which are waiting for stream flow control
    /// credits, and executes the given function on each `Stream`
    ///
    /// The `stream::Controller` will be notified of streams that have been
    /// closed to allow for further streams to be opened.
    pub fn iterate_stream_flow_credits_list<F>(
        &mut self,
        controller: &mut stream::Controller,
        mut func: F,
    ) where
        F: FnMut(&mut S) -> StreamContainerIterationResult,
    {
        iterate_interruptible!(
            self,
            waiting_for_stream_flow_control_credits,
            waiting_for_stream_flow_control_credits_link,
            controller,
            func
        );
    }

    /// Iterates over all `Stream`s which are waiting for transmission,
    /// and executes the given function on each `Stream`
    ///
    /// The `stream::Controller` will be notified of streams that have been
    /// closed to allow for further streams to be opened.
    pub fn iterate_transmission_list<F>(&mut self, controller: &mut stream::Controller, mut func: F)
    where
        F: FnMut(&mut S) -> StreamContainerIterationResult,
    {
        iterate_interruptible!(
            self,
            waiting_for_transmission,
            waiting_for_transmission_link,
            controller,
            func
        );
    }

    /// Iterates over all `Stream`s which are waiting for retransmission,
    /// and executes the given function on each `Stream`
    ///
    /// The `stream::Controller` will be notified of streams that have been
    /// closed to allow for further streams to be opened.
    pub fn iterate_retransmission_list<F>(
        &mut self,
        controller: &mut stream::Controller,
        mut func: F,
    ) where
        F: FnMut(&mut S) -> StreamContainerIterationResult,
    {
        iterate_interruptible!(
            self,
            waiting_for_retransmission,
            waiting_for_retransmission_link,
            controller,
            func
        );
    }

    /// Iterates over all `Stream`s which are part of this container, and executes
    /// the given function on each `Stream`
    ///
    /// The `stream::Controller` will be notified of streams that have been
    /// closed to allow for further streams to be opened.
    pub fn iterate_streams<F>(&mut self, controller: &mut stream::Controller, mut func: F)
    where
        F: FnMut(&mut S),
    {
        // Note: We can not use iterate_uninterruptible here, because that
        // iteration will extract nodes from the list but not automatically insert
        // all Nodes back into the main interest list. `stream_map` is not
        // populated by the interest maps.

        for stream in self.stream_map.iter() {
            debug_assert!(stream.tree_link.is_linked());

            let mut mut_stream = stream.inner.borrow_mut();
            func(&mut *mut_stream);
            let interests = mut_stream.get_stream_interests();

            // Update the interest lists here
            // Safety: The stream reference is obtained from the RBTree, which
            // stores it's nodes as `Rc`
            let stream_node_rc = unsafe { stream_node_rc_from_ref(stream) };
            self.interest_lists.update_interests(
                &stream_node_rc,
                interests,
                StreamContainerIterationResult::Continue,
            );
        }

        if !self.interest_lists.done_streams.is_empty() {
            // Cleanup all `done` streams after we finished interacting with all
            // of them.
            self.finalize_done_streams(controller);
        }
    }

    /// Returns whether or not streams have data to send
    pub fn has_pending_streams(&self) -> bool {
        !self.interest_lists.waiting_for_transmission.is_empty()
            || !self.interest_lists.waiting_for_retransmission.is_empty()
    }
}

impl<S: StreamTrait> timer::Provider for StreamContainer<S> {
    #[inline]
    fn timers<Q: timer::Query>(&self, query: &mut Q) -> timer::Result {
        // TODO denormalize this into a single value
        for stream in self
            .interest_lists
            .waiting_for_stream_flow_control_credits
            .iter()
        {
            stream.inner.borrow().timers(query)?;
        }
        Ok(())
    }
}

impl<S: StreamTrait> transmission::interest::Provider for StreamContainer<S> {
    #[inline]
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        if !self.interest_lists.waiting_for_retransmission.is_empty() {
            query.on_lost_data()?;
        } else if !self.interest_lists.waiting_for_transmission.is_empty() {
            query.on_new_data()?;
        }

        Ok(())
    }
}

/// Return values for iterations over a `Stream` list.
/// The value instructs the iterator whether iteration will be continued.
#[derive(Clone, Copy, Debug)]
pub enum StreamContainerIterationResult {
    /// Continue iteration over the list
    Continue,
    /// Aborts the iteration over a list and add the remaining items at the
    /// back of the list
    BreakAndInsertAtBack,
}


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

use core::task::Waker;

/// Events that are triggered through a state-update on a Stream.
/// E.g. retrieving new data can unblock readers, which is signalled via this
/// Event.
#[derive(Debug, Default)]
pub struct StreamEvents {
    /// A task that is waiting on the reading half of the Stream that might have
    /// been woken up due to an action.
    pub read_wake: Option<Waker>,
    /// A task that is waiting on the writing half of the Stream that might have
    /// been woken up due to an action.
    pub write_wake: Option<Waker>,
}

impl StreamEvents {
    /// Creates a new [`StreamEvents`] instance with no pending events
    pub fn new() -> StreamEvents {
        Default::default()
    }

    /// Returns the amount of [`Waker`] instances that are stored in the event
    pub fn waker_count(&self) -> usize {
        let mut result = 0;
        if self.read_wake.is_some() {
            result += 1;
        }
        if self.write_wake.is_some() {
            result += 1;
        }
        result
    }

    /// Stores a [`Waker`] which was waiting on the reading half of the Stream
    /// to get ready. The [`Waker`] will get woken up by the consumer of
    /// [`StreamEvents`].
    pub fn store_read_waker(&mut self, waker: Waker) {
        self.read_wake = Some(waker);
    }

    /// Stores a [`Waker`] which was waiting on the writing half of the Stream
    /// to get ready. The [`Waker`] will get woken up by the consumer of
    /// [`StreamEvents`].
    pub fn store_write_waker(&mut self, waker: Waker) {
        self.write_wake = Some(waker);
    }

    /// Wakes all [`Waker`] instances that have been stored via `add_waker`.
    /// Waking the [`Waker`]s will remove them from the [`StreamEvents`]s struct.
    pub fn wake_all(&mut self) {
        if let Some(waker) = self.read_wake.take() {
            waker.wake();
        }
        if let Some(waker) = self.write_wake.take() {
            waker.wake();
        }
    }
}


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

use crate::{
    contexts::{OnTransmitError, WriteContext},
    stream::{
        incoming_connection_flow_controller::IncomingConnectionFlowController,
        outgoing_connection_flow_controller::OutgoingConnectionFlowController,
        receive_stream::ReceiveStream,
        send_stream::SendStream,
        stream_events::StreamEvents,
        stream_interests::{StreamInterestProvider, StreamInterests},
        StreamError,
    },
};
use core::{task::Context, time::Duration};
use s2n_quic_core::{
    ack, endpoint,
    frame::{stream::StreamRef, MaxStreamData, ResetStream, StopSending, StreamDataBlocked},
    stream::{ops, StreamId},
    time::{timer, Timestamp},
    transport,
    varint::VarInt,
};

/// Configuration values for a Stream
#[derive(Debug)]
pub struct StreamConfig {
    /// The [`Stream`]s identifier
    pub stream_id: StreamId,
    /// The type of the local endpoint
    pub local_endpoint_type: endpoint::Type,
    /// The connection-wide flow controller for receiving data
    pub incoming_connection_flow_controller: IncomingConnectionFlowController,
    /// The connection-wide flow controller for sending data
    pub outgoing_connection_flow_controller: OutgoingConnectionFlowController,
    /// The initial flow control window for receiving data
    pub initial_receive_window: VarInt,
    /// The desired flow control window that we want to maintain on the receiving side
    pub desired_flow_control_window: u32,
    /// The initial flow control window for sending data
    pub initial_send_window: VarInt,
    /// The maximum buffered amount of data on the sending side
    pub max_send_buffer_size: u32,
}

/// A trait which represents an internally used `Stream`
pub trait StreamTrait: StreamInterestProvider + timer::Provider + core::fmt::Debug {
    /// Creates a new `Stream` instance with the given configuration
    fn new(config: StreamConfig) -> Self;

    /// Returns the Streams ID
    fn stream_id(&self) -> StreamId;

    // These functions are called from the packet delivery thread

    /// This is called when a `STREAM_DATA` frame had been received for
    /// this stream
    fn on_data(
        &mut self,
        frame: &StreamRef,
        events: &mut StreamEvents,
    ) -> Result<(), transport::Error>;

    /// This is called when a `STREAM_DATA_BLOCKED` frame had been received for
    /// this stream
    fn on_stream_data_blocked(
        &mut self,
        frame: &StreamDataBlocked,
        events: &mut StreamEvents,
    ) -> Result<(), transport::Error>;

    /// This is called when a `RESET_STREAM` frame had been received for
    /// this stream
    fn on_reset(
        &mut self,
        frame: &ResetStream,
        events: &mut StreamEvents,
    ) -> Result<(), transport::Error>;

    /// This is called when a `MAX_STREAM_DATA` frame had been received for
    /// this stream
    fn on_max_stream_data(
        &mut self,
        frame: &MaxStreamData,
        events: &mut StreamEvents,
    ) -> Result<(), transport::Error>;

    /// This is called when a `STOP_SENDING` frame had been received for
    /// this stream
    fn on_stop_sending(
        &mut self,
        frame: &StopSending,
        events: &mut StreamEvents,
    ) -> Result<(), transport::Error>;

    /// This method gets called when a packet delivery got acknowledged
    fn on_packet_ack<A: ack::Set>(&mut self, ack_set: &A, events: &mut StreamEvents);

    /// This method gets called when a packet loss is reported
    fn on_packet_loss<A: ack::Set>(&mut self, ack_set: &A, events: &mut StreamEvents);

    /// Updates the period at which `STREAM_DATA_BLOCKED` frames are sent to the peer
    /// if the application is blocked by peer limits.
    fn update_blocked_sync_period(&mut self, blocked_sync_period: Duration);

    /// Called when the connection timer expires
    fn on_timeout(&mut self, now: Timestamp);

    /// This method gets called when a stream gets reset due to a reason that is
    /// not related to a frame. E.g. due to a connection failure.
    fn on_internal_reset(&mut self, error: StreamError, events: &mut StreamEvents);

    /// This method is called when the application drops the connection
    ///
    /// The stream should finish any pending operations and close
    fn on_flush(&mut self, error: StreamError, events: &mut StreamEvents);

    /// Queries the component for any outgoing frames that need to get sent
    fn on_transmit<W: WriteContext>(&mut self, context: &mut W) -> Result<(), OnTransmitError>;

    /// This method is called when a connection window is available
    fn on_connection_window_available(&mut self);

    // These functions are called from the client API

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

/// The implementation of a `Stream`.
/// This is mostly a facade over the reading and writing half of the `Stream`.
#[derive(Debug)]
pub struct StreamImpl {
    /// The stream ID
    pub(super) stream_id: StreamId,
    /// Manages the receiving side of the stream
    pub(super) receive_stream: ReceiveStream,
    /// Set to `true` when this stream has a sending side
    has_send: bool,
    /// Manages the sending side of the stream
    pub(super) send_stream: SendStream,
}

impl StreamImpl {
    fn poll_request_impl(
        &mut self,
        request: &mut ops::Request,
        context: Option<&Context>,
    ) -> Result<ops::Response, StreamError> {
        let mut response = ops::Response::default();
        if let Some(rx) = request.rx.as_mut() {
            match self.receive_stream.poll_request(rx, context) {
                Ok(rx) => response.rx = Some(rx),
                Err(err) => {
                    // only fail the request if tx is not set
                    if response.tx.is_none() {
                        return Err(err);
                    } else {
                        response.rx = Some(ops::rx::Response {
                            status: ops::Status::Reset(err),
                            ..Default::default()
                        });
                    }
                }
            }
        }

        if let Some(tx) = request.tx.as_mut() {
            match self.send_stream.poll_request(tx, context) {
                Ok(tx) => response.tx = Some(tx),
                Err(err) => {
                    // only fail the request if rx is not set
                    if response.rx.is_none() {
                        return Err(err);
                    } else {
                        response.tx = Some(ops::tx::Response {
                            status: ops::Status::Reset(err),
                            ..Default::default()
                        });
                    }
                }
            }
        }

        Ok(response)
    }
}

impl StreamTrait for StreamImpl {
    fn new(config: StreamConfig) -> StreamImpl {
        let receive_is_closed = config.stream_id.stream_type().is_unidirectional()
            && config.stream_id.initiator() == config.local_endpoint_type;
        let send_is_closed = config.stream_id.stream_type().is_unidirectional()
            && config.stream_id.initiator() != config.local_endpoint_type;

        StreamImpl {
            stream_id: config.stream_id,
            receive_stream: ReceiveStream::new(
                receive_is_closed,
                config.incoming_connection_flow_controller,
                config.initial_receive_window,
                config.desired_flow_control_window,
            ),
            has_send: !send_is_closed,
            send_stream: SendStream::new(
                config.outgoing_connection_flow_controller,
                send_is_closed,
                config.initial_send_window,
                config.max_send_buffer_size,
            ),
        }
    }

    #[inline]
    fn stream_id(&self) -> StreamId {
        self.stream_id
    }

    // These functions are called from the packet delivery thread

    #[inline]
    fn on_data(
        &mut self,
        frame: &StreamRef,
        events: &mut StreamEvents,
    ) -> Result<(), transport::Error> {
        self.receive_stream.on_data(frame, events)
    }

    #[inline]
    fn on_stream_data_blocked(
        &mut self,
        frame: &StreamDataBlocked,
        events: &mut StreamEvents,
    ) -> Result<(), transport::Error> {
        self.receive_stream.on_stream_data_blocked(frame, events)
    }

    #[inline]
    fn on_reset(
        &mut self,
        frame: &ResetStream,
        events: &mut StreamEvents,
    ) -> Result<(), transport::Error> {
        self.receive_stream.on_reset(frame, events)
    }

    #[inline]
    fn on_max_stream_data(
        &mut self,
        frame: &MaxStreamData,
        events: &mut StreamEvents,
    ) -> Result<(), transport::Error> {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-19.10
        //# An endpoint that
        //# receives a MAX_STREAM_DATA frame for a receive-only stream MUST
        //# terminate the connection with error STREAM_STATE_ERROR.
        if !self.has_send {
            return Err(transport::Error::STREAM_STATE_ERROR
                .with_reason("MAX_STREAM_DATA sent on receive-only stream"));
        }

        self.send_stream.on_max_stream_data(frame, events)
    }

    #[inline]
    fn on_stop_sending(
        &mut self,
        frame: &StopSending,
        events: &mut StreamEvents,
    ) -> Result<(), transport::Error> {
        self.send_stream.on_stop_sending(frame, events)
    }

    #[inline]
    fn on_packet_ack<A: ack::Set>(&mut self, ack_set: &A, events: &mut StreamEvents) {
        self.receive_stream.on_packet_ack(ack_set);
        self.send_stream.on_packet_ack(ack_set, events);
    }

    #[inline]
    fn on_packet_loss<A: ack::Set>(&mut self, ack_set: &A, _events: &mut StreamEvents) {
        self.receive_stream.on_packet_loss(ack_set);
        self.send_stream.on_packet_loss(ack_set);
    }

    #[inline]
    fn update_blocked_sync_period(&mut self, blocked_sync_period: Duration) {
        self.send_stream
            .update_blocked_sync_period(blocked_sync_period);
    }

    #[inline]
    fn on_timeout(&mut self, now: Timestamp) {
        self.send_stream.on_timeout(now)
    }

    #[inline]
    fn on_internal_reset(&mut self, error: StreamError, events: &mut StreamEvents) {
        self.receive_stream.on_internal_reset(error, events);
        self.send_stream.on_internal_reset(error, events);
    }

    #[inline]
    fn on_flush(&mut self, error: StreamError, events: &mut StreamEvents) {
        // flushing a receive stream is the same as resetting it
        self.receive_stream.on_internal_reset(error, events);

        // tell the send stream to flush any pending data
        self.send_stream.on_flush(error, events);
    }

    #[inline]
    fn on_transmit<W: WriteContext>(&mut self, context: &mut W) -> Result<(), OnTransmitError> {
        // Query the receiving side for outgoing data
        self.receive_stream.on_transmit(self.stream_id, context)?;
        // And the sending side
        self.send_stream.on_transmit(self.stream_id, context)
    }

    #[inline]
    fn on_connection_window_available(&mut self) {
        self.send_stream.on_connection_window_available()
    }

    // These functions are called from the client API

    fn poll_request(
        &mut self,
        request: &mut ops::Request,
        context: Option<&Context>,
    ) -> Result<ops::Response, StreamError> {
        #[cfg(debug_assertions)]
        let contract: crate::stream::contract::Request = (&*request).into();

        let result = self.poll_request_impl(request, context);

        #[cfg(debug_assertions)]
        contract.validate_response(request, result.as_ref(), context);

        result
    }
}

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

impl StreamInterestProvider for StreamImpl {
    #[inline]
    fn stream_interests(&self, interests: &mut StreamInterests) {
        self.send_stream.stream_interests(interests);
        self.receive_stream.stream_interests(interests);
    }
}


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

//! A collection of a all the interactions a `Stream` is interested in

use crate::transmission::interest::{Interest, Query, QueryBreak, Result};

/// A collection of a all the interactions a `Stream` is interested in
#[derive(Debug, Default, Copy, Clone, PartialEq, Eq)]
pub struct StreamInterests {
    /// Is `true` if the `Stream` wants to transmit data but is blocked on
    /// insufficient connection flow control credits
    pub connection_flow_control_credits: bool,
    /// Is `true` if the `Stream` wants to transmit data but is blocked on
    /// insufficient stream flow control credits
    pub stream_flow_control_credits: bool,
    /// Is `true` if the `Stream` is still wanting to make progress. Otherwise
    /// the stream will be removed from the `Stream` map.
    pub retained: bool,
    /// Is `true` if the component is interested in packet acknowledge and
    /// loss information
    pub delivery_notifications: bool,
    /// Transmission interest for the component
    pub transmission: Interest,
}

impl StreamInterests {
    #[inline]
    pub fn merge(&mut self, other: &Self) {
        self.connection_flow_control_credits |= other.connection_flow_control_credits;
        self.stream_flow_control_credits |= other.stream_flow_control_credits;
        self.retained |= other.retained;
        self.delivery_notifications |= other.delivery_notifications;
        let _ = self.transmission.on_interest(other.transmission);
    }

    #[inline]
    pub fn with_transmission<F: FnOnce(&mut TransmissionInterest) -> Result>(&mut self, f: F) {
        let mut interest = TransmissionInterest(&mut self.transmission);
        let _ = f(&mut interest);
    }
}

pub struct TransmissionInterest<'a>(&'a mut Interest);

impl<'a> Query for TransmissionInterest<'a> {
    fn on_interest(&mut self, interest: Interest) -> Result {
        debug_assert_ne!(
            interest,
            Interest::Forced,
            "streams are not allowed to force transmission"
        );

        match (*self.0, interest) {
            // we don't need to keep querying if we're already at the max interest
            (Interest::LostData, _) | (_, Interest::LostData) => {
                *self.0 = Interest::LostData;
                return Err(QueryBreak);
            }
            (Interest::NewData, _) | (_, Interest::NewData) => *self.0 = Interest::NewData,
            _ => {}
        }

        Ok(())
    }
}

/// A type which can provide it's Stream interests
pub trait StreamInterestProvider {
    /// Returns all interactions the object is interested in.
    fn stream_interests(&self, interests: &mut StreamInterests);

    fn get_stream_interests(&self) -> StreamInterests {
        let mut interests = StreamInterests::default();
        self.stream_interests(&mut interests);
        interests
    }
}


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

pub use crate::contexts::testing::*;
use crate::{
    stream::{
        incoming_connection_flow_controller::IncomingConnectionFlowController,
        outgoing_connection_flow_controller::OutgoingConnectionFlowController,
        stream_impl::StreamConfig, stream_interests::StreamInterests, StreamEvents, StreamImpl,
        StreamTrait,
    },
    transmission,
};
use bytes::Bytes;
use core::task::{Context, Poll, Waker};
use futures_test::task::{new_count_waker, AwokenCount};
use s2n_quic_core::{
    application::Error as ApplicationErrorCode,
    endpoint,
    frame::{stream::Stream as StreamFrame, Frame, ResetStream, StreamDataBlocked},
    packet::number::{PacketNumber, PacketNumberSpace},
    stream::{ops, StreamError, StreamId, StreamType},
    time::{clock::testing as time, Timestamp},
    transport,
    varint::VarInt,
};

#[test]
fn idle_stream_does_not_write_data() {
    let mut test_env = setup_stream_test_env();
    test_env.assert_write_frames(0);
}

macro_rules! assert_matches {
    ($a:expr, $b:pat $(,)?) => {
        match $a {
            $b => {}
            ref value => {
                panic!("value {:?} did not match {}", value, stringify!($b))
            }
        }
    };
}
pub(super) use assert_matches;

/// Creates a `STREAM_DATA` frame
pub fn stream_data<Data>(
    stream_id: StreamId,
    offset: VarInt,
    data: Data,
    is_fin: bool,
) -> StreamFrame<Data> {
    StreamFrame {
        offset,
        data,
        stream_id: stream_id.into(),
        is_last_frame: false,
        is_fin,
    }
}

/// Asserts that a `Result` type contains a TransportError with the given
/// error code.
pub fn assert_is_transport_error<T: core::fmt::Debug>(
    result: Result<T, transport::Error>,
    expected: transport::Error,
) {
    let actual = result.unwrap_err();
    assert_eq!(expected.code, actual.code);
}

/// Generates test data using a pattern which is identifieable. For a given
/// offset in the Stream the utilized data will always be the same. This allows
/// us to do some simple validation checking whether a receiver received the
/// expected data without exactly knowing the actual sent data.
pub fn gen_pattern_test_data(offset: VarInt, len: usize) -> Vec<u8> {
    let mut data = Vec::new();
    data.reserve(len);

    fn data_for_offset(offset: u64) -> u8 {
        (offset % 256) as u8
    }

    for i in 0..len {
        let current_offset: u64 = Into::<u64>::into(offset) + i as u64;
        data.push(data_for_offset(current_offset));
    }

    data
}

pub fn gen_pattern_test_chunks(mut offset: VarInt, lens: &[usize]) -> Vec<bytes::Bytes> {
    lens.iter()
        .map(|size| {
            let data = bytes::Bytes::from(gen_pattern_test_data(offset, *size));
            offset += *size;
            data
        })
        .collect::<Vec<_>>()
}

/// Defines whether a wakeup is expected.
/// `None` means there are no expectations. `Some(true)` expects a wakeup,
/// `Some(false)` does not.
#[derive(Debug, Copy, Clone, PartialEq)]
pub struct ExpectWakeup(pub Option<bool>);

/// Creates an application space packet number with the given value
pub fn pn(nr: usize) -> PacketNumber {
    PacketNumberSpace::ApplicationData.new_packet_number(VarInt::new(nr as u64).unwrap())
}

/// Creates Stream Interests from an array of strings
///
/// The following interests are supported:
/// - ack => delivery_notifications
/// - tx => transmission::Interest::NewData
/// - fin => finalization
/// - cf => connection_flow_control_credits
pub fn stream_interests(interests: &[&str]) -> StreamInterests {
    let mut result = StreamInterests {
        retained: true,
        ..Default::default()
    };
    for interest in interests {
        match *interest {
            "ack" => result.delivery_notifications = true,
            "tx" => result.transmission = transmission::Interest::NewData,
            "lost" => result.transmission = transmission::Interest::LostData,
            "fin" => result.retained = false,
            "cf" => result.connection_flow_control_credits = true,
            "sf" => result.stream_flow_control_credits = true,
            other => unreachable!("Unsupported interest {}", other),
        }
    }
    result
}

/// Creates Transmission Interests from an array of strings
///
/// The following interests are supported:
/// - tx => transmission::Interest::NewData
pub fn transmission_interests(interests: &[&str]) -> transmission::Interest {
    let mut result = transmission::Interest::default();
    for interest in interests {
        match *interest {
            "tx" => result = transmission::Interest::NewData,
            "lost" => result = transmission::Interest::LostData,
            other => unreachable!("Unsupported interest {}", other),
        }
    }
    result
}

/// Holds a set of associated objects that act as a test environment for a single
/// [`StreamImpl`].
pub struct TestEnvironment {
    pub sent_frames: OutgoingFrameBuffer,
    pub stream: StreamImpl,
    pub rx_connection_flow_controller: IncomingConnectionFlowController,
    pub tx_connection_flow_controller: OutgoingConnectionFlowController,
    pub wake_counter: AwokenCount,
    pub waker: Waker,
    pub current_time: Timestamp,
    pub transmission_constraint: transmission::Constraint,
    pub endpoint: endpoint::Type,
}

impl TestEnvironment {
    // These are the defaults for configuration values which are utilized for
    // most tests.
    // In order to test that config values are not accidentally mixed up in the
    // library, we use slightly different values for those. The exact numbers
    // should not matter too much - higher numbers will require tests to take
    // longer.

    pub const DEFAULT_INITIAL_CONNECTION_RECEIVE_WINDOW: u64 = 100 * 1024;
    pub const DEFAULT_INITIAL_CONNECTION_SEND_WINDOW: u64 = 100 * 1024;
    pub const DEFAULT_INITIAL_RECEIVE_WINDOW: u64 = 4096;
    pub const DEFAULT_INITIAL_SEND_WINDOW: u64 = 8 * 1024;
    pub const DEFAULT_MAX_SEND_BUFFER_SIZE: usize = 16 * 1024;

    /// Asserts that the given byte array can be read from the stream
    pub fn assert_receive_data(&mut self, data: &'static [u8]) {
        assert_eq!(
            Poll::Ready(Ok(Some(Bytes::from_static(data)))),
            self.poll_pop(),
            "the stream should receive the correct bytes"
        );
    }

    /// Asserts that no data is available for reading and that the stream
    /// has not yet finished
    pub fn assert_no_read_data(&mut self) {
        assert_eq!(
            Poll::Pending,
            self.poll_pop(),
            "the stream should return Poll::Pending"
        );
    }

    /// Asserts that the stream has reached end of stream
    pub fn assert_end_of_stream(&mut self) {
        assert_eq!(
            Poll::Ready(Ok(None)),
            self.poll_pop(),
            "the stream should have finished receiving data"
        );
    }

    /// Asserts that the returns an error when pop is called
    pub fn assert_pop_error(&mut self) {
        let poll_result = self.poll_pop();

        match poll_result {
            Poll::Ready(Err(_)) => {}
            _ => panic!("Expected the stream to be reset"),
        }
    }

    /// Feeds `amount` of bytes into the stream at the given offset
    pub fn feed_data(&mut self, mut offset: VarInt, amount: usize) {
        let mut remaining = amount;
        while remaining > 0 {
            let to_write = core::cmp::min(remaining, 4096);
            let data = vec![0u8; to_write];
            let mut events = StreamEvents::new();
            assert_eq!(
                Ok(()),
                self.stream.on_data(
                    &stream_data(self.stream.stream_id, offset, &data[..], false),
                    &mut events
                )
            );
            offset += VarInt::from_u32(to_write as u32);
            remaining -= to_write;
            events.wake_all();
        }
    }

    /// Consumes all currently available data from the stream
    pub fn consume_all_data(&mut self) -> usize {
        let mut result = 0;
        loop {
            let poll_result = self.poll_pop();

            match poll_result {
                Poll::Pending => break, // Consumed all data
                Poll::Ready(Ok(Some(data))) => {
                    result += data.len();
                }
                Poll::Ready(Ok(None)) => break, // Consumed all data to end of stream
                _ => panic!("Unexpected read result {:?}", poll_result),
            }
        }

        result
    }

    /// Queries the stream for pending outgoing frames.
    /// Asserts that `expected_frames` had been written.
    /// The frames will get written into `sent_frames`.
    pub fn assert_write_frames(&mut self, expected_frames: usize) {
        let prev_written = self.sent_frames.len();
        let mut write_ctx = MockWriteContext::new(
            self.current_time,
            &mut self.sent_frames,
            self.transmission_constraint,
            transmission::Mode::Normal,
            self.endpoint,
        );
        assert!(self
            .rx_connection_flow_controller
            .on_transmit(&mut write_ctx)
            .is_ok());
        let _ = self.stream.on_transmit(&mut write_ctx);
        self.sent_frames.flush();
        assert_eq!(
            prev_written + expected_frames,
            self.sent_frames.len(),
            "Unexpected amount of written frames"
        );
    }

    /// Ensures a single frame was written
    pub fn transmit(&mut self) -> Option<WrittenFrame> {
        let mut write_ctx = MockWriteContext::new(
            self.current_time,
            &mut self.sent_frames,
            self.transmission_constraint,
            transmission::Mode::Normal,
            self.endpoint,
        );
        self.stream.on_transmit(&mut write_ctx).ok()?;
        self.sent_frames.flush();

        self.sent_frames.pop_front()
    }

    /// Asserts that a stream data frame for a given offset and size was emitted
    pub fn assert_write_of(
        &mut self,
        expected_offset: VarInt,
        expected_size: usize,
        expect_eof: bool,
        expect_is_last_frame: bool,
        expected_packet_number: PacketNumber,
    ) {
        let mut sent_frame = self.transmit().expect("no frame was written");
        assert_eq!(
            expected_packet_number, sent_frame.packet_nr,
            "packet number mismatch"
        );

        let decoded_frame = sent_frame.as_frame();
        // These assertions are on individual fields to see more easily
        // where things are failing.
        if let Frame::Stream(mut stream_frame) = decoded_frame {
            assert_eq!(expected_offset, stream_frame.offset, "offset mismatch");
            assert_eq!(expect_eof, stream_frame.is_fin, "FIN mismatch");
            assert_eq!(expected_size, stream_frame.data.len(), "size mismatch");
            assert_eq!(
                expect_is_last_frame, stream_frame.is_last_frame,
                "is_last_frame mismatch"
            );
            assert_eq!(
                gen_pattern_test_data(expected_offset, expected_size),
                stream_frame.data.as_less_safe_slice_mut(),
                "data mismatch",
            );
        } else {
            panic!("Expected a Stream frame, but got {:?}", decoded_frame);
        }
    }

    /// Asserts that a RESET frame was transmitted
    pub fn assert_write_reset_frame(
        &mut self,
        expected_error_code: ApplicationErrorCode,
        expected_packet_number: PacketNumber,
        expected_final_size: VarInt,
    ) {
        let mut sent_frame = self.transmit().expect("no frame was written");
        assert_eq!(
            expected_packet_number, sent_frame.packet_nr,
            "packet number mismatch"
        );

        assert_eq!(
            Frame::ResetStream(ResetStream {
                stream_id: self.stream.stream_id.into(),
                application_error_code: expected_error_code.into(),
                final_size: expected_final_size,
            }),
            sent_frame.as_frame()
        );
    }

    /// Asserts that a STREAM_DATA_BLOCKED frame was transmitted
    pub fn assert_write_stream_data_blocked_frame(
        &mut self,
        expected_stream_data_limit: VarInt,
        expected_packet_number: PacketNumber,
    ) {
        let mut sent_frame = self.transmit().expect("no frame was written");
        assert_eq!(
            expected_packet_number, sent_frame.packet_nr,
            "packet number mismatch"
        );

        assert_eq!(
            Frame::StreamDataBlocked(StreamDataBlocked {
                stream_id: self.stream.stream_id.into(),
                stream_data_limit: expected_stream_data_limit,
            }),
            sent_frame.as_frame()
        );
    }

    /// Acknowledges a packet with a given packet number as received.
    /// `expect_writer_wakeup` specifies whether we expect the wake counter to
    /// get increased due to this operation.
    pub fn ack_packet(&mut self, packet_number: PacketNumber, expect_writer_wakeup: ExpectWakeup) {
        let old_wake_count = self.wake_counter.get();
        self.rx_connection_flow_controller
            .on_packet_ack(&packet_number);
        let mut events = StreamEvents::new();
        self.stream.on_packet_ack(&packet_number, &mut events);
        events.wake_all();
        let new_wake_count = self.wake_counter.get();
        let was_woken = new_wake_count > old_wake_count;
        if let ExpectWakeup(Some(wakeup_expected)) = expect_writer_wakeup {
            assert_eq!(wakeup_expected, was_woken, "Unexpected wakeup through ACK");
        }
    }

    /// Declares a packet with a given packet number as lost
    pub fn nack_packet(&mut self, packet_number: PacketNumber) {
        self.rx_connection_flow_controller
            .on_packet_loss(&packet_number);
        let mut events = StreamEvents::new();
        self.stream.on_packet_loss(&packet_number, &mut events);
    }

    pub fn run_request(
        &mut self,
        request: &mut ops::Request,
        with_context: bool,
    ) -> Result<ops::Response, StreamError> {
        let context = if with_context {
            Some(Context::from_waker(&self.waker))
        } else {
            None
        };
        self.stream.poll_request(request, context.as_ref())
    }

    pub fn poll_request(
        &mut self,
        request: &mut ops::Request,
    ) -> Poll<Result<ops::Response, StreamError>> {
        match self.run_request(request, true)?.into_poll() {
            Poll::Pending => Poll::Pending,
            Poll::Ready(response) => Poll::Ready(Ok(response)),
        }
    }

    pub fn poll_push(&mut self, chunk: Bytes) -> Poll<Result<(), StreamError>> {
        match self
            .run_request(ops::Request::default().send(&mut [chunk]), true)?
            .into_poll()
        {
            Poll::Pending => Poll::Pending,
            Poll::Ready(_) => Poll::Ready(Ok(())),
        }
    }

    pub fn poll_finish(&mut self) -> Poll<Result<(), StreamError>> {
        match self
            .run_request(ops::Request::default().finish().flush(), true)?
            .into_poll()
        {
            Poll::Pending => Poll::Pending,
            Poll::Ready(_) => Poll::Ready(Ok(())),
        }
    }

    pub fn poll_pop(&mut self) -> Poll<Result<Option<Bytes>, StreamError>> {
        let mut chunks = [Bytes::new()];
        match self
            .run_request(ops::Request::default().receive(&mut chunks), true)?
            .into_poll()
        {
            Poll::Pending => Poll::Pending,
            Poll::Ready(response) => {
                let chunk = if response.rx().expect("invalid response").chunks.consumed == 1 {
                    Some(core::mem::replace(&mut chunks[0], Bytes::new()))
                } else {
                    None
                };
                Poll::Ready(Ok(chunk))
            }
        }
    }

    pub fn reset(&mut self, error_code: ApplicationErrorCode) -> Result<(), StreamError> {
        self.run_request(ops::Request::default().reset(error_code), false)?;
        Ok(())
    }

    pub fn stop_sending(&mut self, error_code: ApplicationErrorCode) -> Result<(), StreamError> {
        self.run_request(ops::Request::default().stop_sending(error_code), false)?;
        Ok(())
    }
}

#[derive(Debug, Copy, Clone, PartialEq)]
pub struct TestEnvironmentConfig {
    pub stream_id: StreamId,
    pub initial_receive_window: u64,
    pub desired_flow_control_window: u32,
    pub initial_send_window: u64,
    pub initial_connection_send_window_size: u64,
    pub initial_connection_receive_window_size: u64,
    pub desired_connection_flow_control_window: u32,
    pub max_send_buffer_size: usize,
    pub transmission_constraint: transmission::Constraint,
    pub local_endpoint_type: endpoint::Type,
    pub max_packet_size: Option<usize>,
}

impl Default for TestEnvironmentConfig {
    fn default() -> Self {
        Self::new(endpoint::Type::Server)
    }
}

impl TestEnvironmentConfig {
    pub fn new(local_endpoint_type: endpoint::Type) -> Self {
        TestEnvironmentConfig {
            local_endpoint_type,
            stream_id: StreamId::initial(
                local_endpoint_type.peer_type(),
                StreamType::Bidirectional,
            ),
            initial_receive_window: TestEnvironment::DEFAULT_INITIAL_RECEIVE_WINDOW,
            desired_flow_control_window: TestEnvironment::DEFAULT_INITIAL_RECEIVE_WINDOW as u32,
            initial_send_window: TestEnvironment::DEFAULT_INITIAL_SEND_WINDOW,
            initial_connection_send_window_size:
                TestEnvironment::DEFAULT_INITIAL_CONNECTION_SEND_WINDOW,
            initial_connection_receive_window_size:
                TestEnvironment::DEFAULT_INITIAL_CONNECTION_RECEIVE_WINDOW,
            desired_connection_flow_control_window:
                TestEnvironment::DEFAULT_INITIAL_CONNECTION_RECEIVE_WINDOW as u32,
            max_send_buffer_size: TestEnvironment::DEFAULT_MAX_SEND_BUFFER_SIZE,
            transmission_constraint: transmission::Constraint::None,
            max_packet_size: None,
        }
    }
}

/// Sets up a test environment for Stream testing with default parameters
pub fn setup_stream_test_env() -> TestEnvironment {
    let config = TestEnvironmentConfig::default();
    setup_stream_test_env_with_config(config)
}

/// Sets up a test environment for Stream testing with custom parameters
pub fn setup_stream_test_env_with_config(config: TestEnvironmentConfig) -> TestEnvironment {
    let rx_connection_flow_controller = IncomingConnectionFlowController::new(
        VarInt::new(config.initial_connection_receive_window_size).unwrap(),
        config.desired_connection_flow_control_window,
    );

    let tx_connection_flow_controller = OutgoingConnectionFlowController::new(
        VarInt::new(config.initial_connection_send_window_size).unwrap(),
    );

    let stream = StreamImpl::new(StreamConfig {
        incoming_connection_flow_controller: rx_connection_flow_controller.clone(),
        outgoing_connection_flow_controller: tx_connection_flow_controller.clone(),
        local_endpoint_type: config.local_endpoint_type,
        stream_id: config.stream_id,
        initial_receive_window: VarInt::new(config.initial_receive_window).unwrap(),
        desired_flow_control_window: config.desired_flow_control_window,
        initial_send_window: VarInt::new(config.initial_send_window).unwrap(),
        max_send_buffer_size: config.max_send_buffer_size as u32,
    });

    let (waker, wake_counter) = new_count_waker();

    let mut sent_frames = OutgoingFrameBuffer::new();

    sent_frames.set_max_packet_size(config.max_packet_size);

    TestEnvironment {
        sent_frames,
        stream,
        rx_connection_flow_controller,
        tx_connection_flow_controller,
        wake_counter,
        waker,
        current_time: time::now(),
        transmission_constraint: config.transmission_constraint,
        endpoint: config.local_endpoint_type,
    }
}


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

use std::ops::RangeInclusive;

use super::*;
use bolero::{check, generator::*};
use futures_test::task::new_count_waker;
use hashbrown::HashSet;
use s2n_quic_core::varint::VarInt;

#[derive(Debug)]
struct Oracle {
    local_endpoint_type: endpoint::Type,
    initial_local_limits: InitialFlowControlLimits,
    initial_remote_limits: InitialFlowControlLimits,

    max_remote_bidi_opened_nth_idx: Option<u64>,
    max_remote_uni_opened_nth_idx: Option<u64>,
    max_local_bidi_opened_nth_idx: Option<u64>,
    max_local_uni_opened_nth_idx: Option<u64>,

    remote_bidi_open_idx_set: HashSet<u64>,
    remote_uni_open_idx_set: HashSet<u64>,
    local_bidi_open_idx_set: HashSet<u64>,
    local_uni_open_idx_set: HashSet<u64>,
}

impl Oracle {
    fn can_open_remote(&self, stream_type: StreamType, nth_idx: u64) -> bool {
        // the count is +1 since streams are 0-indexed
        let nth_cnt = nth_idx + 1;
        let limit = self.remote_limit(stream_type);

        nth_cnt <= limit
    }

    fn on_open_stream(
        &mut self,
        stream_initiator: endpoint::Type,
        stream_type: StreamType,
        nth_idx: u64,
    ) {
        match (stream_initiator == self.local_endpoint_type, stream_type) {
            (true, StreamType::Bidirectional) => self.max_local_bidi_opened_nth_idx = Some(nth_idx),
            (true, StreamType::Unidirectional) => self.max_local_uni_opened_nth_idx = Some(nth_idx),
            (false, StreamType::Bidirectional) => {
                self.max_remote_bidi_opened_nth_idx = Some(nth_idx)
            }
            (false, StreamType::Unidirectional) => {
                self.max_remote_uni_opened_nth_idx = Some(nth_idx)
            }
        };

        match (stream_initiator == self.local_endpoint_type, stream_type) {
            (true, StreamType::Bidirectional) => self.local_bidi_open_idx_set.insert(nth_idx),
            (true, StreamType::Unidirectional) => self.local_uni_open_idx_set.insert(nth_idx),
            (false, StreamType::Bidirectional) => self.remote_bidi_open_idx_set.insert(nth_idx),
            (false, StreamType::Unidirectional) => self.remote_uni_open_idx_set.insert(nth_idx),
        };
    }

    // Returns the range of stream ids which can be opened.
    //
    // None is returned if no streams can be opened. None is returned if
    // max_opened_nth_idx stream is > than the nth_idx stream requested. This
    // is because opening a higher stream also opens the lower streams and we
    // do not handle opening a closed stream.
    fn open_stream_range(
        &self,
        stream_initiator: endpoint::Type,
        stream_type: StreamType,
        nth_idx: u64,
    ) -> Option<RangeInclusive<u64>> {
        let max_opened_nth_idx = match (stream_initiator == self.local_endpoint_type, stream_type) {
            (true, StreamType::Bidirectional) => self.max_local_bidi_opened_nth_idx,
            (true, StreamType::Unidirectional) => self.max_local_uni_opened_nth_idx,
            (false, StreamType::Bidirectional) => self.max_remote_bidi_opened_nth_idx,
            (false, StreamType::Unidirectional) => self.max_remote_uni_opened_nth_idx,
        };

        let stream_nth_idx_iter = if let Some(max_opened_nth_idx) = max_opened_nth_idx {
            // idx already opened.. return
            if max_opened_nth_idx >= nth_idx {
                return None;
            }

            // +1 to get the next stream to open
            max_opened_nth_idx + 1..=nth_idx
        } else {
            0..=nth_idx
        };

        Some(stream_nth_idx_iter)
    }

    fn can_close(
        &self,
        stream_initiator: endpoint::Type,
        stream_type: StreamType,
        nth_idx: u64,
    ) -> bool {
        match (stream_initiator == self.local_endpoint_type, stream_type) {
            (true, StreamType::Bidirectional) => self.local_bidi_open_idx_set.contains(&nth_idx),
            (true, StreamType::Unidirectional) => self.local_uni_open_idx_set.contains(&nth_idx),
            (false, StreamType::Bidirectional) => self.remote_bidi_open_idx_set.contains(&nth_idx),
            (false, StreamType::Unidirectional) => self.remote_uni_open_idx_set.contains(&nth_idx),
        }
    }

    fn on_close_stream(
        &mut self,
        stream_initiator: endpoint::Type,
        stream_type: StreamType,
        nth_idx: u64,
    ) {
        match (stream_initiator == self.local_endpoint_type, stream_type) {
            (true, StreamType::Bidirectional) => {
                self.local_bidi_open_idx_set.take(&nth_idx).unwrap();
            }
            (true, StreamType::Unidirectional) => {
                self.local_uni_open_idx_set.take(&nth_idx).unwrap();
            }
            (false, StreamType::Bidirectional) => {
                self.remote_bidi_open_idx_set.take(&nth_idx).unwrap();
                self.initial_local_limits
                    .max_open_remote_bidirectional_streams += 1;
            }
            (false, StreamType::Unidirectional) => {
                self.remote_uni_open_idx_set.take(&nth_idx).unwrap();
                self.initial_local_limits
                    .max_open_remote_unidirectional_streams += 1;
            }
        };
    }

    fn remote_limit(&self, stream_type: StreamType) -> u64 {
        match stream_type {
            StreamType::Bidirectional => {
                self.initial_local_limits
                    .max_open_remote_bidirectional_streams
            }
            StreamType::Unidirectional => {
                self.initial_local_limits
                    .max_open_remote_unidirectional_streams
            }
        }
        .as_u64()
    }

    fn open_streams_count(&self, stream_initiator: endpoint::Type, stream_type: StreamType) -> u64 {
        match (stream_initiator == self.local_endpoint_type, stream_type) {
            (true, StreamType::Bidirectional) => self.local_bidi_open_idx_set.len() as u64,
            (true, StreamType::Unidirectional) => self.local_uni_open_idx_set.len() as u64,
            (false, StreamType::Bidirectional) => self.remote_bidi_open_idx_set.len() as u64,
            (false, StreamType::Unidirectional) => self.remote_uni_open_idx_set.len() as u64,
        }
    }

    fn on_max_stream_local(&mut self, maximum_streams: u8, direction: StreamDirection) {
        match direction {
            StreamDirection::LocalInitiatedBidirectional => {
                if self
                    .initial_remote_limits
                    .max_open_remote_bidirectional_streams
                    .as_u64()
                    < maximum_streams.into()
                {
                    self.initial_remote_limits
                        .max_open_remote_bidirectional_streams = VarInt::from_u8(maximum_streams);
                }
            }
            StreamDirection::LocalInitiatedUnidirectional => {
                if self
                    .initial_remote_limits
                    .max_open_remote_unidirectional_streams
                    .as_u64()
                    < maximum_streams.into()
                {
                    self.initial_remote_limits
                        .max_open_remote_unidirectional_streams = VarInt::from_u8(maximum_streams);
                }
            }
            StreamDirection::RemoteInitiatedBidirectional
            | StreamDirection::RemoteInitiatedUnidirectional => {
                panic!("should only be called for local endpoints")
            }
        };
    }
}

#[derive(Debug)]
struct Model {
    oracle: Oracle,
    subject: Controller,
}

impl Model {
    fn new(local_endpoint_type: endpoint::Type, limits: Limits) -> Self {
        let (initial_local_limits, initial_remote_limits, stream_limits) =
            limits.as_controller_limits();

        Model {
            oracle: Oracle {
                local_endpoint_type,
                initial_local_limits,
                initial_remote_limits,
                max_remote_bidi_opened_nth_idx: None,
                max_remote_uni_opened_nth_idx: None,
                max_local_bidi_opened_nth_idx: None,
                max_local_uni_opened_nth_idx: None,
                remote_bidi_open_idx_set: HashSet::new(),
                remote_uni_open_idx_set: HashSet::new(),
                local_bidi_open_idx_set: HashSet::new(),
                local_uni_open_idx_set: HashSet::new(),
            },
            subject: Controller::new(
                local_endpoint_type,
                initial_remote_limits,
                initial_local_limits,
                stream_limits,
            ),
        }
    }

    pub fn apply(&mut self, operation: &Operation) {
        match operation {
            Operation::OpenRemoteBidi { nth_idx } => {
                self.on_open_remote(*nth_idx as u64, StreamType::Bidirectional)
            }
            Operation::OpenRemoteUni { nth_idx } => {
                self.on_open_remote(*nth_idx as u64, StreamType::Unidirectional)
            }
            Operation::OpenLocalBidi { nth_idx } => {
                self.on_open_local(*nth_idx as u64, StreamType::Bidirectional)
            }
            Operation::OpenLocalUni { nth_idx } => {
                self.on_open_local(*nth_idx as u64, StreamType::Unidirectional)
            }
            Operation::CloseRemoteBidi { nth_idx } => {
                self.on_close_remote(*nth_idx as u64, StreamType::Bidirectional)
            }
            Operation::CloseRemoteUni { nth_idx } => {
                self.on_close_remote(*nth_idx as u64, StreamType::Unidirectional)
            }
            Operation::CloseLocalBidi { nth_idx } => {
                self.on_close_local(*nth_idx as u64, StreamType::Bidirectional)
            }
            Operation::CloseLocalUni { nth_idx } => {
                self.on_close_local(*nth_idx as u64, StreamType::Unidirectional)
            }
            Operation::MaxStreamLocalBidi { maximum_streams } => self.on_max_stream_local(
                *maximum_streams,
                StreamDirection::LocalInitiatedBidirectional,
            ),
            Operation::MaxStreamLocalUni { maximum_streams } => self.on_max_stream_local(
                *maximum_streams,
                StreamDirection::LocalInitiatedUnidirectional,
            ),
        }
    }

    fn on_max_stream_local(&mut self, maximum_streams: u8, direction: StreamDirection) {
        let stream_type = match direction {
            StreamDirection::LocalInitiatedBidirectional => StreamType::Bidirectional,
            StreamDirection::LocalInitiatedUnidirectional => StreamType::Unidirectional,
            StreamDirection::RemoteInitiatedBidirectional
            | StreamDirection::RemoteInitiatedUnidirectional => panic!("dont expect"),
        };
        let frame = MaxStreams {
            stream_type,
            maximum_streams: VarInt::from_u32(maximum_streams.try_into().unwrap()),
        };
        self.subject.on_max_streams(&frame);
        self.oracle.on_max_stream_local(maximum_streams, direction);
    }

    fn on_open_local(&mut self, nth_idx: u64, stream_type: StreamType) {
        let (waker, _wake_counter) = new_count_waker();
        let mut token = connection::OpenToken::new();

        let stream_initiator = self.oracle.local_endpoint_type;

        //-------------
        let stream_nth_idx_iter =
            match self
                .oracle
                .open_stream_range(stream_initiator, stream_type, nth_idx)
            {
                Some(val) => val,
                None => return,
            };

        for stream_nth_idx in stream_nth_idx_iter {
            let can_open = self.can_open_local(stream_type);

            let res = self.subject.poll_open_local_stream(
                stream_type,
                &mut token,
                &Context::from_waker(&waker),
            );

            if can_open {
                assert!(res.is_ready());
                self.oracle
                    .on_open_stream(stream_initiator, stream_type, stream_nth_idx);
            } else {
                assert!(res.is_pending())
            }
        }
    }

    fn on_open_remote(&mut self, nth_idx: u64, stream_type: StreamType) {
        let stream_initiator = self.oracle.local_endpoint_type.peer_type();

        //-------------
        let stream_nth_idx_iter =
            match self
                .oracle
                .open_stream_range(stream_initiator, stream_type, nth_idx)
            {
                Some(val) => val,
                None => return,
            };

        let start_stream =
            StreamId::nth(stream_initiator, stream_type, *stream_nth_idx_iter.start()).unwrap();
        let end_stream =
            StreamId::nth(stream_initiator, stream_type, *stream_nth_idx_iter.end()).unwrap();

        let stream_iter = StreamIter::new(start_stream, end_stream);
        let res = self.subject.on_open_remote_stream(stream_iter);

        if self.can_open_remote(stream_type, nth_idx) {
            for stream_nth_idx in stream_nth_idx_iter {
                self.oracle
                    .on_open_stream(stream_initiator, stream_type, stream_nth_idx);
            }
            res.unwrap();
        } else {
            res.expect_err("limits violated");
        }
    }

    fn on_close_local(&mut self, nth_idx: u64, stream_type: StreamType) {
        let stream_initiator = self.oracle.local_endpoint_type;

        //-------------
        if !self
            .oracle
            .can_close(stream_initiator, stream_type, nth_idx)
        {
            return;
        }

        self.oracle
            .on_close_stream(stream_initiator, stream_type, nth_idx);
        let stream_id = StreamId::nth(stream_initiator, stream_type, nth_idx).unwrap();
        self.subject.on_close_stream(stream_id);
    }

    fn on_close_remote(&mut self, nth_idx: u64, stream_type: StreamType) {
        let stream_initiator = self.oracle.local_endpoint_type.peer_type();

        //-------------
        if !self
            .oracle
            .can_close(stream_initiator, stream_type, nth_idx)
        {
            return;
        }

        self.oracle
            .on_close_stream(stream_initiator, stream_type, nth_idx);
        let stream_id = StreamId::nth(stream_initiator, stream_type, nth_idx).unwrap();
        self.subject.on_close_stream(stream_id);
    }

    fn can_open_local(&self, stream_type: StreamType) -> bool {
        let available_stream_capacity = self
            .subject
            .available_local_initiated_stream_capacity(stream_type);

        available_stream_capacity > VarInt::from_u8(0)
    }

    fn can_open_remote(&self, stream_type: StreamType, nth_idx: u64) -> bool {
        self.oracle.can_open_remote(stream_type, nth_idx)
    }

    /// Check that the subject and oracle match.
    pub fn invariants(&self) {
        // ---------
        let mut stream_initiator = self.oracle.local_endpoint_type.peer_type();
        let mut stream_type = StreamType::Bidirectional;
        assert_eq!(
            self.subject.remote_bidi_controller.open_stream_count(),
            self.oracle
                .open_streams_count(stream_initiator, stream_type)
        );
        assert_eq!(
            self.subject
                .remote_initiated_max_streams_latest_value(stream_type)
                .as_u64(),
            self.oracle.remote_limit(stream_type)
        );

        // ---------
        stream_initiator = self.oracle.local_endpoint_type.peer_type();
        stream_type = StreamType::Unidirectional;
        assert_eq!(
            self.subject.remote_uni_controller.open_stream_count(),
            self.oracle
                .open_streams_count(stream_initiator, stream_type)
        );
        assert_eq!(
            self.subject
                .remote_initiated_max_streams_latest_value(stream_type)
                .as_u64(),
            self.oracle.remote_limit(stream_type)
        );

        // ---------
        stream_initiator = self.oracle.local_endpoint_type;
        stream_type = StreamType::Bidirectional;
        assert_eq!(
            self.subject.local_bidi_controller.open_stream_count(),
            self.oracle
                .open_streams_count(stream_initiator, stream_type)
        );
        self.subject.local_bidi_controller.check_integrity();

        // ---------
        stream_initiator = self.oracle.local_endpoint_type;
        stream_type = StreamType::Unidirectional;
        assert_eq!(
            self.subject.local_uni_controller.open_stream_count(),
            self.oracle
                .open_streams_count(stream_initiator, stream_type)
        );
        self.subject.local_uni_controller.check_integrity();
    }
}

#[derive(Debug, TypeGenerator)]
enum Operation {
    // max_local_limit: max_remote_uni_stream (initial_local_limits)
    // transmit: max_streams
    OpenRemoteBidi { nth_idx: u8 },
    CloseRemoteBidi { nth_idx: u8 },

    // max_local_limit: max_remote_uni_stream (initial_local_limits)
    // transmit: max_streams
    OpenRemoteUni { nth_idx: u8 },
    CloseRemoteUni { nth_idx: u8 },

    // max_local_limit: max_local_bidi_stream
    // peer_stream_limit: peer_max_bidi_stream (initial_remote_limits)
    //
    // limits: max_local_bidi_stream.min(peer_max_bidi_stream)
    // transmit: streams_blocked
    OpenLocalBidi { nth_idx: u8 },
    CloseLocalBidi { nth_idx: u8 },
    MaxStreamLocalBidi { maximum_streams: u8 },

    // max_local_limit: max_local_uni_stream
    // peer_stream_limit: peer_max_uni_stream (initial_remote_limits)
    //
    // limits: max_local_uni_stream.min(peer_max_uni_stream)
    // transmit: streams_blocked
    OpenLocalUni { nth_idx: u8 },
    CloseLocalUni { nth_idx: u8 },
    MaxStreamLocalUni { maximum_streams: u8 },
}

#[derive(Debug, TypeGenerator, Clone, Copy)]
struct Limits {
    // OpenRemoteBidi (initial_local_limits)
    initial_local_max_remote_bidi: u8,

    // OpenRemoteUni (initial_local_limits)
    initial_local_max_remote_uni: u8,

    // OpenLocalBidi (initial_remote_limits)
    //  initial_remote_max_local_bidi.min(app_max_local_bidi)
    initial_remote_max_local_bidi: u8,
    app_max_local_bidi: u8,

    // OpenLocalUni (initial_remote_limits)
    //  initial_remote_max_local_uni.min(app_max_local_uni)
    initial_remote_max_local_uni: u8,
    app_max_local_uni: u8,
}

impl Limits {
    fn as_controller_limits(
        &self,
    ) -> (
        InitialFlowControlLimits,
        InitialFlowControlLimits,
        stream::Limits,
    ) {
        let mut initial_local_limits = InitialFlowControlLimits::default();
        let mut initial_remote_limits = InitialFlowControlLimits::default();
        let stream_limits = stream::Limits {
            max_open_local_unidirectional_streams: (self.app_max_local_uni as u64)
                .try_into()
                .unwrap(),
            max_open_local_bidirectional_streams: (self.app_max_local_bidi as u64)
                .try_into()
                .unwrap(),
            ..Default::default()
        };

        // OpenRemoteBidi (initial_local_limits)
        initial_local_limits.max_open_remote_bidirectional_streams =
            VarInt::from_u32(self.initial_local_max_remote_bidi.into());

        // OpenRemoteUni (initial_local_limits)
        initial_local_limits.max_open_remote_unidirectional_streams =
            VarInt::from_u32(self.initial_local_max_remote_uni.into());

        // OpenLocalBidi (initial_remote_limits)
        //  initial_remote_max_local_bidi.min(app_max_local_bidi)
        initial_remote_limits.max_open_remote_bidirectional_streams =
            VarInt::from_u32(self.initial_remote_max_local_bidi.into());

        // OpenLocalUni (initial_remote_limits)
        //  initial_remote_max_local_uni.min(app_max_local_uni)
        initial_remote_limits.max_open_remote_unidirectional_streams =
            VarInt::from_u32(self.initial_remote_max_local_uni.into());

        (initial_local_limits, initial_remote_limits, stream_limits)
    }
}

#[test]
fn model_test() {
    check!()
        .with_type::<(Limits, Vec<Operation>)>()
        .for_each(|(limits, operations)| {
            let local_endpoint_type = endpoint::Type::Server;

            let mut model = Model::new(local_endpoint_type, *limits);
            for operation in operations.iter() {
                model.apply(operation);
            }

            model.invariants();
        })
}


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

use crate::{
    connection::open_token,
    contexts::OnTransmitError,
    sync::{OnceSync, PeriodicSync, ValueToFrameWriter},
    transmission,
    transmission::WriteContext,
};
use core::{
    task::{Context, Poll, Waker},
    time::Duration,
};
use s2n_quic_core::{
    ack,
    frame::{self, MaxStreams, StreamsBlocked},
    packet::number::PacketNumber,
    stream::{limits::LocalLimits, StreamId},
    time::{timer, Timestamp},
    varint::VarInt,
};
use smallvec::SmallVec;

// The amount of wakers that may be tracked before allocating to the heap.
const WAKERS_INITIAL_CAPACITY: usize = 5;

/// The LocalInitiated controller controls streams initiated locally
#[derive(Debug)]
pub(super) struct LocalInitiated<L: LocalLimits, OpenNotify: OpenNotifyBehavior> {
    /// The max stream limit specified by the local endpoint.
    ///
    /// Used to restrict the number of concurrent streams the local
    /// connection can open.
    max_local_limit: L,
    /// The cumulative stream limit specified by the remote endpoint.
    ///
    /// Can be updated when MAX_STREAMS frame is received.
    peer_cumulative_stream_limit: VarInt,
    wakers: SmallVec<[Waker; WAKERS_INITIAL_CAPACITY]>,
    streams_blocked_sync: PeriodicSync<VarInt, StreamsBlockedToFrameWriter>,
    /// opened_streams is needed to track the latest opened stream since
    /// peer_stream_limit is a cumulative limit.
    opened_streams: VarInt,
    closed_streams: VarInt,
    /// Keeps track of all of the issued open tokens
    token_counter: open_token::Counter,
    /// Keeps track of all of the expired open tokens
    expired_token: open_token::Token,
    open_notify: OpenNotify,
}

impl<L: LocalLimits, OpenNotify: OpenNotifyBehavior> LocalInitiated<L, OpenNotify> {
    pub fn new(initial_peer_maximum_streams: VarInt, max_local_limit: L) -> Self {
        Self {
            max_local_limit,
            peer_cumulative_stream_limit: initial_peer_maximum_streams,
            wakers: SmallVec::new(),
            streams_blocked_sync: PeriodicSync::new(),
            opened_streams: VarInt::from_u8(0),
            closed_streams: VarInt::from_u8(0),
            token_counter: open_token::Counter::new(),
            expired_token: open_token::Token::new(),
            open_notify: Default::default(),
        }
    }

    pub fn on_max_streams(&mut self, frame: &MaxStreams) {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-4.6
        //# MAX_STREAMS frames that do not increase the stream limit MUST be ignored.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-19.11
        //# MAX_STREAMS frames that do not increase the stream limit MUST be ignored.
        if self.peer_cumulative_stream_limit >= frame.maximum_streams {
            return;
        }

        self.peer_cumulative_stream_limit = frame.maximum_streams;

        // We now have more capacity from the peer so stop sending STREAMS_BLOCKED frames
        self.streams_blocked_sync.stop_sync();

        self.wake_unblocked();
    }

    pub fn update_sync_period(&mut self, blocked_sync_period: Duration) {
        self.streams_blocked_sync
            .update_sync_period(blocked_sync_period);
    }

    #[inline]
    pub fn poll_open_stream(
        &mut self,
        open_token: &mut open_token::Token,
        context: &Context,
    ) -> Poll<()> {
        if self.available_stream_capacity() < VarInt::from_u32(1) {
            if let Some(index) = open_token.index(&self.expired_token) {
                let prev = &self.wakers[index];
                // update the waker if it's changed
                if !prev.will_wake(context.waker()) {
                    self.wakers[index] = context.waker().clone();
                }
            } else {
                // Store a waker that can be woken when we get more credit
                self.wakers.push(context.waker().clone());
                // give them a waker to remember their position in the list
                *open_token = self.token_counter.next();
            }

            //= https://www.rfc-editor.org/rfc/rfc9000#section-4.6
            //# An endpoint that is unable to open a new stream due to the peer's
            //# limits SHOULD send a STREAMS_BLOCKED frame (Section 19.14).

            //= https://www.rfc-editor.org/rfc/rfc9000#section-19.14
            //# A sender SHOULD send a STREAMS_BLOCKED frame (type=0x16 or 0x17) when
            //# it wishes to open a stream but is unable to do so due to the maximum
            //# stream limit set by its peer; see Section 19.11.
            if self.peer_capacity() < VarInt::from_u32(1) {
                self.streams_blocked_sync
                    .request_delivery(self.peer_cumulative_stream_limit)
            }

            return Poll::Pending;
        }

        // reset the open token since they're no longer blocked
        open_token.clear();

        Poll::Ready(())
    }

    #[inline]
    pub fn on_open_stream(&mut self) {
        self.opened_streams += 1;
        self.open_notify.on_open_stream();

        self.check_integrity();
    }

    pub fn on_close_stream(&mut self) {
        self.closed_streams += 1;

        self.wake_unblocked();
        self.check_integrity();
    }

    /// The number of streams that may be opened by the local application, respecting both
    /// the local concurrent streams limit and the peer's stream limits.
    #[inline]
    pub fn available_stream_capacity(&self) -> VarInt {
        let local_capacity = self
            .max_local_limit
            .as_varint()
            .saturating_sub(self.open_stream_count());
        local_capacity.min(self.peer_capacity())
    }

    /// The current number of streams that can be opened according to the peer's limits
    #[inline]
    fn peer_capacity(&self) -> VarInt {
        self.peer_cumulative_stream_limit
            .saturating_sub(self.opened_streams)
    }

    /// Wake all wakers
    fn wake_all(&mut self) {
        self.wakers
            .drain(..self.wakers.len())
            .for_each(|waker| waker.wake())
    }

    /// Wakes the wakers that have been unblocked by the current amount
    /// of available local stream capacity.
    fn wake_unblocked(&mut self) {
        let unblocked_wakers_count = self
            .wakers
            .len()
            .min(self.available_stream_capacity().as_u64() as usize);

        self.wakers
            .drain(..unblocked_wakers_count)
            .for_each(|waker| waker.wake());

        // keep track of the number of tokens that have expired
        self.expired_token.expire(unblocked_wakers_count);
    }

    /// Returns the number of streams currently open
    #[inline]
    pub fn open_stream_count(&self) -> VarInt {
        self.opened_streams - self.closed_streams
    }

    #[inline]
    pub fn total_open_stream_count(&self) -> VarInt {
        self.opened_streams
    }

    #[inline]
    pub fn on_timeout(&mut self, now: Timestamp) {
        self.streams_blocked_sync.on_timeout(now);
    }

    #[inline]
    pub fn on_packet_ack<A: ack::Set>(&mut self, ack_set: &A) {
        self.open_notify.on_packet_ack(ack_set);
        self.streams_blocked_sync.on_packet_ack(ack_set)
    }

    #[inline]
    pub fn on_packet_loss<A: ack::Set>(&mut self, ack_set: &A) {
        self.open_notify.on_packet_loss(ack_set);
        self.streams_blocked_sync.on_packet_loss(ack_set)
    }

    #[inline]
    pub fn on_transmit<W: WriteContext>(
        &mut self,
        stream_id: StreamId,
        context: &mut W,
    ) -> Result<(), OnTransmitError> {
        if context.ack_elicitation().is_ack_eliciting() && self.streams_blocked_sync.has_delivered()
        {
            // We are already sending an ack-eliciting packet, so no need to send another STREAMS_BLOCKED.
            // This matches the RFC requirement below for STREAM_DATA_BLOCKED and DATA_BLOCKED.
            //= https://www.rfc-editor.org/rfc/rfc9000#section-4.1
            //# To keep the
            //# connection from closing, a sender that is flow control limited SHOULD
            //# periodically send a STREAM_DATA_BLOCKED or DATA_BLOCKED frame when it
            //# has no ack-eliciting packets in flight.
            self.streams_blocked_sync
                .skip_delivery(context.current_time());
        } else {
            self.streams_blocked_sync.on_transmit(stream_id, context)?;
        }

        self.open_notify.on_transmit(stream_id, context)?;

        Ok(())
    }

    pub fn close(&mut self) {
        self.wake_all();
        self.streams_blocked_sync.stop_sync();
        self.open_notify.close();
    }

    #[inline]
    pub fn check_integrity(&self) {
        if cfg!(debug_assertions) {
            assert!(
                self.closed_streams <= self.opened_streams,
                "Cannot close more streams than previously opened"
            );
            assert!(
                self.open_stream_count() <= self.max_local_limit.as_varint(),
                "Cannot have more outgoing streams open concurrently than
                the max_local_limit"
            );
        }
    }
}

impl<L: LocalLimits, OpenNotify: OpenNotifyBehavior> timer::Provider
    for LocalInitiated<L, OpenNotify>
{
    #[inline]
    fn timers<Q: timer::Query>(&self, query: &mut Q) -> timer::Result {
        self.streams_blocked_sync.timers(query)?;
        Ok(())
    }
}

impl<L: LocalLimits, OpenNotify: OpenNotifyBehavior> transmission::interest::Provider
    for LocalInitiated<L, OpenNotify>
{
    #[inline]
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        self.streams_blocked_sync.transmission_interest(query)?;
        self.open_notify.transmission_interest(query)?;
        Ok(())
    }
}

/// Writes the `STREAMS_BLOCKED` frames.
#[derive(Debug, Default)]
pub(super) struct StreamsBlockedToFrameWriter {}

impl ValueToFrameWriter<VarInt> for StreamsBlockedToFrameWriter {
    #[inline]
    fn write_value_as_frame<W: WriteContext>(
        &self,
        value: VarInt,
        stream_id: StreamId,
        context: &mut W,
    ) -> Option<PacketNumber> {
        context.write_frame(&StreamsBlocked {
            stream_type: stream_id.stream_type(),
            stream_limit: value,
        })
    }
}

/// Interface for determining what to do for notifying the peer about opening
/// the stream.
///
/// By default, the trait provides no-op implementations for all of the methods.
pub trait OpenNotifyBehavior: Default + transmission::interest::Provider {
    #[inline]
    fn on_open_stream(&mut self) {}

    #[inline]
    fn on_transmit<W: WriteContext>(
        &mut self,
        stream_id: StreamId,
        context: &mut W,
    ) -> Result<(), OnTransmitError> {
        let _ = stream_id;
        let _ = context;
        Ok(())
    }

    #[inline]
    fn on_packet_ack<A: ack::Set>(&mut self, ack_set: &A) {
        let _ = ack_set;
    }

    #[inline]
    fn on_packet_loss<A: ack::Set>(&mut self, ack_set: &A) {
        let _ = ack_set;
    }

    #[inline]
    fn close(&mut self) {}
}

/// Defines the open notify behavior for unidirectional streams
///
/// Since locally-initiated, unidirectional streams can only send data, all of
/// these operations are no-op, as the peer will become aware of the stream once
/// the local application starts sending on it.
#[derive(Debug, Default)]
pub(super) struct OpenNotifyUnidirectional;

impl OpenNotifyBehavior for OpenNotifyUnidirectional {}

impl transmission::interest::Provider for OpenNotifyUnidirectional {
    #[inline]
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        _query: &mut Q,
    ) -> transmission::interest::Result {
        Ok(())
    }
}

/// Defines the open notify behavior for bidirectional streams
///
/// This will send an empty STREAM frame to the peer for the largest
/// opened, locally-initiated bidirectional stream. This prevents a deadlock
/// where the local application opens the stream and tries to receive on it.
#[derive(Debug, Default)]
pub(super) struct OpenNotifyBidirectional {
    max_opened: OnceSync<(), OpenNotifyFrameWriter>,
}

impl OpenNotifyBehavior for OpenNotifyBidirectional {
    #[inline]
    fn on_open_stream(&mut self) {
        self.max_opened.force_delivery(())
    }

    #[inline]
    fn on_transmit<W: WriteContext>(
        &mut self,
        stream_id: StreamId,
        context: &mut W,
    ) -> Result<(), OnTransmitError> {
        self.max_opened.on_transmit(stream_id, context)
    }

    #[inline]
    fn on_packet_ack<A: ack::Set>(&mut self, ack_set: &A) {
        let _ = self.max_opened.on_packet_ack(ack_set);
    }

    #[inline]
    fn on_packet_loss<A: ack::Set>(&mut self, ack_set: &A) {
        self.max_opened.on_packet_loss(ack_set)
    }

    #[inline]
    fn close(&mut self) {
        self.max_opened.stop_sync();
    }
}

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

/// Writes an empty STREAM frame to the peer, indicating the stream has been created
#[derive(Debug, Default)]
struct OpenNotifyFrameWriter;

impl ValueToFrameWriter<()> for OpenNotifyFrameWriter {
    #[inline]
    fn write_value_as_frame<W: WriteContext>(
        &self,
        _value: (),
        stream_id: StreamId,
        context: &mut W,
    ) -> Option<PacketNumber> {
        context.write_frame(&frame::Stream {
            stream_id: stream_id.into(),
            is_last_frame: false,
            is_fin: false,
            offset: VarInt::from_u32(0),
            data: &[0u8; 0][..],
        })
    }
}