//! Code for implementing flow control (stream-level).

use std::sync::Arc;

use postage::watch;

use tor_cell::relaycell::flow_ctrl::{Xoff, Xon, XonKbpsEwma};

use tor_cell::relaycell::msg::AnyRelayMsg;

use tor_cell::relaycell::{RelayMsg, UnparsedRelayMsg};

use super::params::FlowCtrlParameters;

use super::window::state::{HalfStreamWindowFlowCtrl, WindowFlowCtrl};

use super::xon_xoff::reader::DrainRateRequest;

#[cfg(feature = "flowctl-cc")]

use super::xon_xoff::state::{HalfStreamXonXoffFlowCtrl, XonXoffFlowCtrl};

use crate::Result;

use crate::congestion::sendme;

use crate::util::notify::NotifySender;

/// Private internals of [`StreamFlowCtrl`].

#[enum_dispatch::enum_dispatch]

#[derive(Debug)]

enum StreamFlowCtrlInner {

    /// "legacy" sendme-window-based flow control.

    WindowBased(WindowFlowCtrl),

    /// XON/XOFF flow control.

    #[cfg(feature = "flowctl-cc")]

    XonXoffBased(XonXoffFlowCtrl),

}

/// Manages the circuit reactor's flow control for a stream.

///

/// Note that the flow control logic can be distributed across multiple parts of Arti.

/// For example some flow control logic will exist in the circuit reactor,

/// but other logic will exist in the stream's `DataStream`.

/// So this doesn't include all flow control logic.

#[derive(Debug)]

pub(crate) struct StreamFlowCtrl {

    /// Private internal enum.

    inner: StreamFlowCtrlInner,

}

impl StreamFlowCtrl {

    /// Returns a new sendme-window-based [`StreamFlowCtrl`].

    /// Returns a new xon/xoff-based [`StreamFlowCtrl`].

    #[cfg(feature = "flowctl-cc")]

    /// Once this end of the stream is closed and the stream becomes a

    /// half-stream (`HalfStream`),

    /// this method will turn the flow control object into a version

    /// that is designed to be used for half-streams.

            }

            #[cfg(feature = "flowctl-cc")]

            }

        };

}

// forward all trait methods to the inner enum

impl FlowCtrlHooks for StreamFlowCtrl {

}

/// Methods that can be called on a [`StreamFlowCtrl`].

///

/// We use a trait so that we can use `enum_dispatch` on the inner [`StreamFlowCtrlInner`] enum.

#[enum_dispatch::enum_dispatch(StreamFlowCtrlInner)]

pub(crate) trait FlowCtrlHooks {

    /// Whether this stream is ready to send `msg`.

    fn can_send<M: RelayMsg>(&self, msg: &M) -> bool;

    /// Inform the flow control code that we're about to send `msg`.

    /// Returns an error if the message should not be sent,

    /// and the circuit should be closed.

    // TODO: Consider having this method wrap the message in a type that

    // "proves" we've applied flow control. This would make it easier to apply

    // flow control earlier, e.g. in `OpenStreamEntStream`, without introducing

    // ambiguity in the sending function as to whether flow control has already

    // been applied or not.

    fn about_to_send(&mut self, msg: &AnyRelayMsg) -> Result<()>;

    /// Handle an incoming sendme.

    ///

    /// On success, return the number of cells left in the window.

    ///

    /// On failure, return an error: the caller should close the stream or

    /// circuit with a protocol error.

    ///

    /// Takes the [`UnparsedRelayMsg`] so that we don't even try to decode it if we're not using the

    /// correct type of flow control.

    fn put_for_incoming_sendme(&mut self, msg: UnparsedRelayMsg) -> Result<()>;

    /// Handle an incoming XON message.

    ///

    /// Takes the [`UnparsedRelayMsg`] so that we don't even try to decode it if we're not using the

    /// correct type of flow control.

    fn handle_incoming_xon(&mut self, msg: UnparsedRelayMsg) -> Result<()>;

    /// Handle an incoming XOFF message.

    ///

    /// Takes the [`UnparsedRelayMsg`] so that we don't even try to decode it if we're not using the

    /// correct type of flow control.

    fn handle_incoming_xoff(&mut self, msg: UnparsedRelayMsg) -> Result<()>;

    /// Check if we should send an XON message.

    ///

    /// If we should, then returns the XON message that should be sent.

    /// Returns an error if XON/XOFF messages aren't supported for this type of flow control.

    fn maybe_send_xon(&mut self, rate: XonKbpsEwma, buffer_len: usize) -> Result<Option<Xon>>;

    /// Check if we should send an XOFF message.

    ///

    /// If we should, then returns the XOFF message that should be sent.

    /// Returns an error if XON/XOFF messages aren't supported for this type of flow control.

    fn maybe_send_xoff(&mut self, buffer_len: usize) -> Result<Option<Xoff>>;

}

/// Manages flow control for a half-stream (`HalfStream`).

#[derive(Debug)]

pub(crate) struct HalfStreamFlowCtrl {

    /// Private internal enum.

    inner: HalfStreamFlowCtrlInner,

}

/// Private internals of [`HalfStreamFlowCtrl`].

#[enum_dispatch::enum_dispatch]

#[derive(Debug)]

enum HalfStreamFlowCtrlInner {

    /// "legacy" sendme-window-based flow control.

    WindowBased(HalfStreamWindowFlowCtrl),

    /// XON/XOFF flow control.

    #[cfg(feature = "flowctl-cc")]

    XonXoffBased(HalfStreamXonXoffFlowCtrl),

}

/// Methods that can be called on a [`HalfStreamFlowCtrl`].

///

/// We use a trait so that we can use `enum_dispatch` on the inner [`HalfStreamFlowCtrlInner`] enum.

/// While this may seem unnecessary since this trait currently only has two methods,

/// it's consistent with the [`FlowCtrlHooks`] trait above.

#[enum_dispatch::enum_dispatch(HalfStreamFlowCtrlInner)]

pub(crate) trait HalfStreamFlowCtrlHooks {

    /// Handle some number of dropped stream messages.

    ///

    /// We don't know what kinds of stream messages were dropped, only the number of them.

    ///

    /// This method exists because currently the stream entry may drop some incoming stream

    /// messages and they would never be processed by this flow control object otherwise.

    fn handle_incoming_dropped(&mut self, msg_count: u16) -> Result<()>;

    /// Handle an incoming message.

    ///

    /// If it's a flow control message, it will be consumed and `None` will be returned.

    /// Otherwise the original message will be returned.

    ///

    /// Takes the [`UnparsedRelayMsg`] so that we don't even try to decode it if we're not using the

    /// correct type of flow control.

    fn handle_incoming_msg(&mut self, msg: UnparsedRelayMsg) -> Result<Option<UnparsedRelayMsg>>;

}

// forward all trait methods to the inner enum

impl HalfStreamFlowCtrlHooks for HalfStreamFlowCtrl {

}

/// A newtype wrapper for a tor stream rate limit that makes the units explicit.

#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]

pub(crate) struct StreamRateLimit {

    /// The rate in bytes/s.

    rate: u64,

}

impl StreamRateLimit {

    /// A maximum rate limit.

    pub(crate) const MAX: Self = Self::new_bytes_per_sec(u64::MAX);

    /// A rate limit of 0.

    pub(crate) const ZERO: Self = Self::new_bytes_per_sec(0);

    /// A new [`StreamRateLimit`] with `rate` bytes/s.

    /// The rate in bytes/s.

}

impl std::fmt::Display for StreamRateLimit {

}