//! Rlpx protocol multiplexer and satellite stream
//!
//! A Satellite is a Stream that primarily drives a single `RLPx` subprotocol but can also handle
//! additional subprotocols.
//!
//! Most of other subprotocols are "dependent satellite" protocols of "eth" and not a fully standalone protocol, for example "snap", See also [snap protocol](https://github.com/ethereum/devp2p/blob/298d7a77c3bf833641579ecbbb5b13f0311eeeea/caps/snap.md?plain=1#L71)
//! Hence it is expected that the primary protocol is "eth" and the additional protocols are
//! "dependent satellite" protocols.

use std::{
    collections::VecDeque,
    fmt,
    future::Future,
    io,
    pin::{pin, Pin},
    sync::Arc,
    task::{ready, Context, Poll},
};

use crate::{
    capability::{SharedCapabilities, SharedCapability, UnsupportedCapabilityError},
    errors::{EthStreamError, P2PStreamError},
    handshake::EthRlpxHandshake,
    p2pstream::DisconnectP2P,
    CanDisconnect, Capability, DisconnectReason, EthStream, P2PStream, UnifiedStatus,
    HANDSHAKE_TIMEOUT,
};
use bytes::{Bytes, BytesMut};
use futures::{Sink, SinkExt, Stream, StreamExt, TryStream, TryStreamExt};
use reth_eth_wire_types::NetworkPrimitives;
use reth_ethereum_forks::ForkFilter;
use tokio::sync::{mpsc, mpsc::UnboundedSender};
use tokio_stream::wrappers::UnboundedReceiverStream;

/// A Stream and Sink type that wraps a raw rlpx stream [`P2PStream`] and handles message ID
/// multiplexing.
#[derive(Debug)]
pub struct RlpxProtocolMultiplexer<St> {
    inner: MultiplexInner<St>,
}

impl<St> RlpxProtocolMultiplexer<St> {
    /// Wraps the raw p2p stream
    pub fn new(conn: P2PStream<St>) -> Self {
        Self {
            inner: MultiplexInner {
                conn,
                protocols: Default::default(),
                out_buffer: Default::default(),
            },
        }
    }

    /// Installs a new protocol on top of the raw p2p stream.
    ///
    /// This accepts a closure that receives a [`ProtocolConnection`] that will yield messages for
    /// the given capability.
    pub fn install_protocol<F, Proto>(
        &mut self,
        cap: &Capability,
        f: F,
    ) -> Result<(), UnsupportedCapabilityError>
    where
        F: FnOnce(ProtocolConnection) -> Proto,
        Proto: Stream<Item = BytesMut> + Send + 'static,
    {
        self.inner.install_protocol(cap, f)
    }

    /// Returns the [`SharedCapabilities`] of the underlying raw p2p stream
    pub const fn shared_capabilities(&self) -> &SharedCapabilities {
        self.inner.shared_capabilities()
    }

    /// Converts this multiplexer into a [`RlpxSatelliteStream`] with the given primary protocol.
    pub fn into_satellite_stream<F, Primary>(
        self,
        cap: &Capability,
        primary: F,
    ) -> Result<RlpxSatelliteStream<St, Primary>, P2PStreamError>
    where
        F: FnOnce(ProtocolProxy) -> Primary,
    {
        let Ok(shared_cap) = self.shared_capabilities().ensure_matching_capability(cap).cloned()
        else {
            return Err(P2PStreamError::CapabilityNotShared)
        };

        let (to_primary, from_wire) = mpsc::unbounded_channel();
        let (to_wire, from_primary) = mpsc::unbounded_channel();
        let proxy = ProtocolProxy {
            shared_cap: shared_cap.clone(),
            from_wire: UnboundedReceiverStream::new(from_wire),
            to_wire,
        };

        let st = primary(proxy);
        Ok(RlpxSatelliteStream {
            inner: self.inner,
            primary: PrimaryProtocol {
                to_primary,
                from_primary: UnboundedReceiverStream::new(from_primary),
                st,
                shared_cap,
            },
            next_outbound: 0,
        })
    }

    /// Converts this multiplexer into a [`RlpxSatelliteStream`] with the given primary protocol.
    ///
    /// Returns an error if the primary protocol is not supported by the remote or the handshake
    /// failed.
    pub async fn into_satellite_stream_with_handshake<F, Fut, Err, Primary>(
        self,
        cap: &Capability,
        handshake: F,
    ) -> Result<RlpxSatelliteStream<St, Primary>, Err>
    where
        F: FnOnce(ProtocolProxy) -> Fut,
        Fut: Future<Output = Result<Primary, Err>>,
        St: Stream<Item = io::Result<BytesMut>> + Sink<Bytes, Error = io::Error> + Unpin,
        P2PStreamError: Into<Err>,
    {
        self.into_satellite_stream_with_tuple_handshake(cap, async move |proxy| {
            let st = handshake(proxy).await?;
            Ok((st, ()))
        })
        .await
        .map(|(st, _)| st)
    }

    /// Converts this multiplexer into a [`RlpxSatelliteStream`] with the given primary protocol.
    ///
    /// Returns an error if the primary protocol is not supported by the remote or the handshake
    /// failed.
    ///
    /// This accepts a closure that does a handshake with the remote peer and returns a tuple of the
    /// primary stream and extra data.
    ///
    /// See also [`UnauthedEthStream::handshake`](crate::UnauthedEthStream)
    pub async fn into_satellite_stream_with_tuple_handshake<F, Fut, Err, Primary, Extra>(
        mut self,
        cap: &Capability,
        handshake: F,
    ) -> Result<(RlpxSatelliteStream<St, Primary>, Extra), Err>
    where
        F: FnOnce(ProtocolProxy) -> Fut,
        Fut: Future<Output = Result<(Primary, Extra), Err>>,
        St: Stream<Item = io::Result<BytesMut>> + Sink<Bytes, Error = io::Error> + Unpin,
        P2PStreamError: Into<Err>,
    {
        let Ok(shared_cap) = self.shared_capabilities().ensure_matching_capability(cap).cloned()
        else {
            return Err(P2PStreamError::CapabilityNotShared.into())
        };

        let (to_primary, from_wire) = mpsc::unbounded_channel();
        let (to_wire, mut from_primary) = mpsc::unbounded_channel();
        let proxy = ProtocolProxy {
            shared_cap: shared_cap.clone(),
            from_wire: UnboundedReceiverStream::new(from_wire),
            to_wire,
        };

        let f = handshake(proxy);
        let mut f = pin!(f);

        // this polls the connection and the primary stream concurrently until the handshake is
        // complete
        loop {
            tokio::select! {
                biased;
                Some(Ok(msg)) = self.inner.conn.next() => {
                    // Ensure the message belongs to the primary protocol
                    let Some(offset) = msg.first().copied()
                    else {
                        return Err(P2PStreamError::EmptyProtocolMessage.into())
                    };
                    if let Some(cap) = self.shared_capabilities().find_by_relative_offset(offset).cloned() {
                            if cap == shared_cap {
                                // delegate to primary
                                let _ = to_primary.send(msg);
                            } else {
                                // delegate to satellite
                                self.inner.delegate_message(&cap, msg);
                            }
                        } else {
                           return Err(P2PStreamError::UnknownReservedMessageId(offset).into())
                        }
                }
                Some(msg) = from_primary.recv() => {
                    self.inner.conn.send(msg).await.map_err(Into::into)?;
                }
                // Poll all subprotocols for new messages
                msg = ProtocolsPoller::new(&mut self.inner.protocols) => {
                     self.inner.conn.send(msg.map_err(Into::into)?).await.map_err(Into::into)?;
                }
                res = &mut f => {
                    let (st, extra) = res?;
                    return Ok((
                        RlpxSatelliteStream {
                            inner: self.inner,
                            primary: PrimaryProtocol {
                                to_primary,
                                from_primary: UnboundedReceiverStream::new(from_primary),
                                st,
                                shared_cap,
                            },
                            next_outbound: 0,
                        },
                        extra,
                    ))
                }
            }
        }
    }

    /// Converts this multiplexer into a [`RlpxSatelliteStream`] with eth protocol as the given
    /// primary protocol and the handshake implementation.
    pub async fn into_eth_satellite_stream<N: NetworkPrimitives>(
        self,
        status: UnifiedStatus,
        fork_filter: ForkFilter,
        handshake: Arc<dyn EthRlpxHandshake>,
        eth_max_message_size: usize,
    ) -> Result<(RlpxSatelliteStream<St, EthStream<ProtocolProxy, N>>, UnifiedStatus), EthStreamError>
    where
        St: Stream<Item = io::Result<BytesMut>> + Sink<Bytes, Error = io::Error> + Unpin,
    {
        let eth_cap = self.inner.conn.shared_capabilities().eth_version()?;
        self.into_satellite_stream_with_tuple_handshake(
            &Capability::eth(eth_cap),
            async move |proxy| {
                let handshake = handshake.clone();
                let mut unauth = UnauthProxy { inner: proxy };
                let their_status = handshake
                    .handshake(&mut unauth, status, fork_filter, HANDSHAKE_TIMEOUT)
                    .await?;
                let eth_stream = EthStream::with_max_message_size(
                    eth_cap,
                    unauth.into_inner(),
                    eth_max_message_size,
                );
                Ok((eth_stream, their_status))
            },
        )
        .await
    }
}

#[derive(Debug)]
struct MultiplexInner<St> {
    /// The raw p2p stream
    conn: P2PStream<St>,
    /// All the subprotocols that are multiplexed on top of the raw p2p stream
    protocols: VecDeque<ProtocolStream>,
    /// Buffer for outgoing messages on the wire.
    out_buffer: OutBuffer,
}

impl<St> MultiplexInner<St> {
    const fn shared_capabilities(&self) -> &SharedCapabilities {
        self.conn.shared_capabilities()
    }

    /// Delegates a message to the matching protocol.
    fn delegate_message(&self, cap: &SharedCapability, msg: BytesMut) -> bool {
        for proto in &self.protocols {
            if proto.shared_cap == *cap {
                proto.send_raw(msg);
                return true
            }
        }
        false
    }

    fn install_protocol<F, Proto>(
        &mut self,
        cap: &Capability,
        f: F,
    ) -> Result<(), UnsupportedCapabilityError>
    where
        F: FnOnce(ProtocolConnection) -> Proto,
        Proto: Stream<Item = BytesMut> + Send + 'static,
    {
        let shared_cap =
            self.conn.shared_capabilities().ensure_matching_capability(cap).cloned()?;
        let (to_satellite, rx) = mpsc::unbounded_channel();
        let proto_conn = ProtocolConnection { from_wire: UnboundedReceiverStream::new(rx) };
        let st = f(proto_conn);
        let st = ProtocolStream { shared_cap, to_satellite, satellite_st: Box::pin(st) };
        self.protocols.push_back(st);
        Ok(())
    }
}

/// Represents a protocol in the multiplexer that is used as the primary protocol.
#[derive(Debug)]
struct PrimaryProtocol<Primary> {
    /// Channel to send messages to the primary protocol.
    to_primary: UnboundedSender<BytesMut>,
    /// Receiver for messages from the primary protocol.
    from_primary: UnboundedReceiverStream<Bytes>,
    /// Shared capability of the primary protocol.
    shared_cap: SharedCapability,
    /// The primary stream.
    st: Primary,
}

/// A Stream and Sink type that acts as a wrapper around a primary `RLPx` subprotocol (e.g. "eth")
///
/// Only emits and sends _non-empty_ messages
#[derive(Debug)]
pub struct ProtocolProxy {
    shared_cap: SharedCapability,
    /// Receives _non-empty_ messages from the wire
    from_wire: UnboundedReceiverStream<BytesMut>,
    /// Sends _non-empty_ messages from the wire
    to_wire: UnboundedSender<Bytes>,
}

impl ProtocolProxy {
    /// Sends a _non-empty_ message on the wire.
    fn try_send(&self, msg: Bytes) -> Result<(), io::Error> {
        if msg.is_empty() {
            // message must not be empty
            return Err(io::ErrorKind::InvalidInput.into())
        }
        self.to_wire.send(self.mask_msg_id(msg)?).map_err(|_| io::ErrorKind::BrokenPipe.into())
    }

    /// Masks the message ID of a message to be sent on the wire.
    #[inline]
    fn mask_msg_id(&self, msg: Bytes) -> Result<Bytes, io::Error> {
        if msg.is_empty() {
            // message must not be empty
            return Err(io::ErrorKind::InvalidInput.into())
        }

        let offset = self.shared_cap.relative_message_id_offset();
        if offset == 0 {
            return Ok(msg);
        }

        let mut masked: BytesMut = msg.into();
        masked[0] = masked[0].checked_add(offset).ok_or(io::ErrorKind::InvalidInput)?;
        Ok(masked.freeze())
    }

    /// Unmasks the message ID of a message received from the wire.
    #[inline]
    fn unmask_id(&self, mut msg: BytesMut) -> Result<BytesMut, io::Error> {
        if msg.is_empty() {
            // message must not be empty
            return Err(io::ErrorKind::InvalidInput.into())
        }
        msg[0] = msg[0]
            .checked_sub(self.shared_cap.relative_message_id_offset())
            .ok_or(io::ErrorKind::InvalidInput)?;
        Ok(msg)
    }
}

impl Stream for ProtocolProxy {
    type Item = Result<BytesMut, io::Error>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let msg = ready!(self.from_wire.poll_next_unpin(cx));
        Poll::Ready(msg.map(|msg| self.get_mut().unmask_id(msg)))
    }
}

impl Sink<Bytes> for ProtocolProxy {
    type Error = io::Error;

    fn poll_ready(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        Poll::Ready(Ok(()))
    }

    fn start_send(self: Pin<&mut Self>, item: Bytes) -> Result<(), Self::Error> {
        self.get_mut().try_send(item)
    }

    fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        Poll::Ready(Ok(()))
    }

    fn poll_close(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        Poll::Ready(Ok(()))
    }
}

impl CanDisconnect<Bytes> for ProtocolProxy {
    fn disconnect(
        &mut self,
        _reason: DisconnectReason,
    ) -> Pin<Box<dyn Future<Output = Result<(), <Self as Sink<Bytes>>::Error>> + Send + '_>> {
        Box::pin(async move { Ok(()) })
    }
}

/// Adapter so the injected `EthRlpxHandshake` can run over a multiplexed `ProtocolProxy`
/// using the same error type expectations (`P2PStreamError`).
#[derive(Debug)]
struct UnauthProxy {
    inner: ProtocolProxy,
}

impl UnauthProxy {
    fn into_inner(self) -> ProtocolProxy {
        self.inner
    }
}

impl Stream for UnauthProxy {
    type Item = Result<BytesMut, P2PStreamError>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.inner.poll_next_unpin(cx).map(|opt| opt.map(|res| res.map_err(P2PStreamError::from)))
    }
}

impl Sink<Bytes> for UnauthProxy {
    type Error = P2PStreamError;

    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready_unpin(cx).map_err(P2PStreamError::from)
    }

    fn start_send(mut self: Pin<&mut Self>, item: Bytes) -> Result<(), Self::Error> {
        self.inner.start_send_unpin(item).map_err(P2PStreamError::from)
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_flush_unpin(cx).map_err(P2PStreamError::from)
    }

    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_close_unpin(cx).map_err(P2PStreamError::from)
    }
}

impl CanDisconnect<Bytes> for UnauthProxy {
    fn disconnect(
        &mut self,
        reason: DisconnectReason,
    ) -> Pin<Box<dyn Future<Output = Result<(), <Self as Sink<Bytes>>::Error>> + Send + '_>> {
        let fut = self.inner.disconnect(reason);
        Box::pin(async move { fut.await.map_err(P2PStreamError::from) })
    }
}

/// A connection channel to receive _`non_empty`_ messages for the negotiated protocol.
///
/// This is a [Stream] that returns raw bytes of the received messages for this protocol.
#[derive(Debug)]
pub struct ProtocolConnection {
    from_wire: UnboundedReceiverStream<BytesMut>,
}

impl Stream for ProtocolConnection {
    type Item = BytesMut;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.from_wire.poll_next_unpin(cx)
    }
}

/// A Stream and Sink type that acts as a wrapper around a primary `RLPx` subprotocol (e.g. "eth")
/// [`EthStream`] and can also handle additional subprotocols.
#[derive(Debug)]
pub struct RlpxSatelliteStream<St, Primary> {
    inner: MultiplexInner<St>,
    primary: PrimaryProtocol<Primary>,
    /// Round-robin cursor for the next outbound producer to poll.
    next_outbound: usize,
}

impl<St, Primary> RlpxSatelliteStream<St, Primary> {
    /// Installs a new protocol on top of the raw p2p stream.
    ///
    /// This accepts a closure that receives a [`ProtocolConnection`] that will yield messages for
    /// the given capability.
    pub fn install_protocol<F, Proto>(
        &mut self,
        cap: &Capability,
        f: F,
    ) -> Result<(), UnsupportedCapabilityError>
    where
        F: FnOnce(ProtocolConnection) -> Proto,
        Proto: Stream<Item = BytesMut> + Send + 'static,
    {
        self.inner.install_protocol(cap, f)
    }

    /// Returns the primary protocol.
    #[inline]
    pub const fn primary(&self) -> &Primary {
        &self.primary.st
    }

    /// Returns mutable access to the primary protocol.
    #[inline]
    pub const fn primary_mut(&mut self) -> &mut Primary {
        &mut self.primary.st
    }

    /// Returns the underlying [`P2PStream`].
    #[inline]
    pub const fn inner(&self) -> &P2PStream<St> {
        &self.inner.conn
    }

    /// Returns mutable access to the underlying [`P2PStream`].
    #[inline]
    pub const fn inner_mut(&mut self) -> &mut P2PStream<St> {
        &mut self.inner.conn
    }

    /// Consumes this type and returns the wrapped [`P2PStream`].
    #[inline]
    pub fn into_inner(self) -> P2PStream<St> {
        self.inner.conn
    }

    /// Polls primary and satellite outbound producers round-robin until the `OutBuffer` is full or
    /// every producer is pending.
    ///
    /// The cursor advances after each producer poll, so a ready producer cannot drain repeatedly
    /// before later producers get a turn.
    fn poll_outbound_producers(&mut self, cx: &mut Context<'_>) -> Result<ProducerPoll, io::Error> {
        let producers = self.inner.protocols.len() + 1;
        let mut pending = 0;

        while pending < producers {
            if self.inner.out_buffer.is_full() {
                return Ok(ProducerPoll::Full)
            }

            if self.next_outbound >= producers {
                self.next_outbound = 0;
            }

            let producer = self.next_outbound;
            self.next_outbound = (self.next_outbound + 1) % producers;

            let msg = if producer == 0 {
                match self.primary.from_primary.poll_next_unpin(cx) {
                    Poll::Ready(Some(msg)) => msg,
                    Poll::Ready(None) => return Ok(ProducerPoll::Closed),
                    Poll::Pending => {
                        pending += 1;
                        continue
                    }
                }
            } else {
                let proto = self
                    .inner
                    .protocols
                    .get_mut(producer - 1)
                    .expect("outbound producer index checked against protocol count");
                match proto.poll_next_unpin(cx) {
                    Poll::Ready(Some(msg)) => msg?,
                    Poll::Ready(None) => return Ok(ProducerPoll::Closed),
                    Poll::Pending => {
                        pending += 1;
                        continue
                    }
                }
            };

            pending = 0;
            self.inner.out_buffer.push_back(msg);
        }

        Ok(ProducerPoll::Pending)
    }
}

impl<St, Primary, PrimaryErr> Stream for RlpxSatelliteStream<St, Primary>
where
    St: Stream<Item = io::Result<BytesMut>> + Sink<Bytes, Error = io::Error> + Unpin,
    Primary: TryStream<Error = PrimaryErr> + Unpin,
    P2PStreamError: Into<PrimaryErr>,
{
    type Item = Result<Primary::Ok, Primary::Error>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let this = self.get_mut();

        loop {
            // first drain the primary stream
            if let Poll::Ready(Some(msg)) = this.primary.st.try_poll_next_unpin(cx) {
                return Poll::Ready(Some(msg))
            }

            let mut conn_ready = true;
            loop {
                match this.inner.conn.poll_ready_unpin(cx) {
                    Poll::Ready(Ok(())) => {
                        if let Some(msg) = this.inner.out_buffer.pop_front() {
                            if let Err(err) = this.inner.conn.start_send_unpin(msg) {
                                return Poll::Ready(Some(Err(err.into())))
                            }
                        } else {
                            break
                        }
                    }
                    Poll::Ready(Err(err)) => {
                        if let Err(disconnect_err) =
                            this.inner.conn.start_disconnect(DisconnectReason::DisconnectRequested)
                        {
                            return Poll::Ready(Some(Err(disconnect_err.into())))
                        }
                        return Poll::Ready(Some(Err(err.into())))
                    }
                    Poll::Pending => {
                        conn_ready = false;
                        break
                    }
                }
            }
            // The connection only buffers frames on `start_send`; `poll_flush` performs the
            // actual writes and flushes the transport once for the batch handed to it above.
            // This also resumes a flush that returned pending on an earlier pass; a no-op if
            // nothing is buffered.
            match this.inner.conn.poll_flush_unpin(cx) {
                Poll::Ready(Ok(())) => {}
                Poll::Ready(Err(err)) => return Poll::Ready(Some(Err(err.into()))),
                Poll::Pending => {
                    conn_ready = false;
                }
            }

            match this.poll_outbound_producers(cx) {
                Ok(ProducerPoll::Pending | ProducerPoll::Full) => {}
                Ok(ProducerPoll::Closed) => return Poll::Ready(None),
                Err(err) => return Poll::Ready(Some(Err(P2PStreamError::Io(err).into()))),
            }

            let mut delegated = false;
            loop {
                // pull messages from connection
                match this.inner.conn.poll_next_unpin(cx) {
                    Poll::Ready(Some(Ok(msg))) => {
                        delegated = true;
                        let Some(offset) = msg.first().copied() else {
                            return Poll::Ready(Some(Err(
                                P2PStreamError::EmptyProtocolMessage.into()
                            )))
                        };
                        // delegate the multiplexed message to the correct protocol
                        if let Some(cap) =
                            this.inner.conn.shared_capabilities().find_by_relative_offset(offset)
                        {
                            if cap == &this.primary.shared_cap {
                                // delegate to primary
                                let _ = this.primary.to_primary.send(msg);
                            } else {
                                // delegate to installed satellite if any
                                for proto in &this.inner.protocols {
                                    if proto.shared_cap == *cap {
                                        proto.send_raw(msg);
                                        break
                                    }
                                }
                            }
                        } else {
                            return Poll::Ready(Some(Err(P2PStreamError::UnknownReservedMessageId(
                                offset,
                            )
                            .into())))
                        }
                    }
                    Poll::Ready(Some(Err(err))) => return Poll::Ready(Some(Err(err.into()))),
                    Poll::Ready(None) => {
                        // connection closed
                        return Poll::Ready(None)
                    }
                    Poll::Pending => break,
                }
            }

            if !conn_ready || (!delegated && this.inner.out_buffer.is_empty()) {
                return Poll::Pending
            }
        }
    }
}

impl<St, Primary, T> Sink<T> for RlpxSatelliteStream<St, Primary>
where
    St: Stream<Item = io::Result<BytesMut>> + Sink<Bytes, Error = io::Error> + Unpin,
    Primary: Sink<T> + Unpin,
    P2PStreamError: Into<<Primary as Sink<T>>::Error>,
{
    type Error = <Primary as Sink<T>>::Error;

    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        let this = self.get_mut();
        if let Err(err) = ready!(this.inner.conn.poll_ready_unpin(cx)) {
            return Poll::Ready(Err(err.into()))
        }
        if let Err(err) = ready!(this.primary.st.poll_ready_unpin(cx)) {
            return Poll::Ready(Err(err))
        }
        Poll::Ready(Ok(()))
    }

    fn start_send(self: Pin<&mut Self>, item: T) -> Result<(), Self::Error> {
        self.get_mut().primary.st.start_send_unpin(item)
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.get_mut().inner.conn.poll_flush_unpin(cx).map_err(Into::into)
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.get_mut().inner.conn.poll_close_unpin(cx).map_err(Into::into)
    }
}

/// Wraps a `RLPx` subprotocol and handles message ID multiplexing.
struct ProtocolStream {
    shared_cap: SharedCapability,
    /// the channel shared with the satellite stream
    to_satellite: UnboundedSender<BytesMut>,
    satellite_st: Pin<Box<dyn Stream<Item = BytesMut> + Send>>,
}

impl ProtocolStream {
    /// Masks the message ID of a message to be sent on the wire.
    #[inline]
    fn mask_msg_id(&self, mut msg: BytesMut) -> Result<Bytes, io::Error> {
        if msg.is_empty() {
            // message must not be empty
            return Err(io::ErrorKind::InvalidInput.into())
        }
        msg[0] = msg[0]
            .checked_add(self.shared_cap.relative_message_id_offset())
            .ok_or(io::ErrorKind::InvalidInput)?;
        Ok(msg.freeze())
    }

    /// Unmasks the message ID of a message received from the wire.
    #[inline]
    fn unmask_id(&self, mut msg: BytesMut) -> Result<BytesMut, io::Error> {
        if msg.is_empty() {
            // message must not be empty
            return Err(io::ErrorKind::InvalidInput.into())
        }
        msg[0] = msg[0]
            .checked_sub(self.shared_cap.relative_message_id_offset())
            .ok_or(io::ErrorKind::InvalidInput)?;
        Ok(msg)
    }

    /// Sends the message to the satellite stream.
    fn send_raw(&self, msg: BytesMut) {
        let _ = self.unmask_id(msg).map(|msg| self.to_satellite.send(msg));
    }
}

impl Stream for ProtocolStream {
    type Item = Result<Bytes, io::Error>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let this = self.get_mut();
        let msg = ready!(this.satellite_st.as_mut().poll_next(cx));
        Poll::Ready(msg.map(|msg| this.mask_msg_id(msg)))
    }
}

impl fmt::Debug for ProtocolStream {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ProtocolStream").field("cap", &self.shared_cap).finish_non_exhaustive()
    }
}

/// Helper to poll multiple protocol streams in a `tokio::select`! branch
struct ProtocolsPoller<'a> {
    protocols: &'a mut VecDeque<ProtocolStream>,
}

impl<'a> ProtocolsPoller<'a> {
    const fn new(protocols: &'a mut VecDeque<ProtocolStream>) -> Self {
        Self { protocols }
    }
}

impl<'a> Future for ProtocolsPoller<'a> {
    type Output = Result<Bytes, P2PStreamError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let protocols = self.protocols.len();
        for _ in 0..protocols {
            let mut proto = self.protocols.pop_front().expect("protocol count checked");
            match proto.poll_next_unpin(cx) {
                Poll::Ready(Some(Err(err))) => {
                    self.protocols.push_back(proto);
                    return Poll::Ready(Err(P2PStreamError::from(err)))
                }
                Poll::Ready(Some(Ok(msg))) => {
                    // Got a message, put protocol back and return the message
                    self.protocols.push_back(proto);
                    return Poll::Ready(Ok(msg));
                }
                _ => {
                    // push it back because we still want to complete the handshake first
                    self.protocols.push_back(proto);
                }
            }
        }

        // All protocols processed, nothing ready
        Poll::Pending
    }
}

/// Soft cap for per-connection outbound `RLPx` messages waiting in the multiplexer.
///
/// The cap is soft because the next message size is only known after polling a protocol stream.
/// The buffer may exceed this by at most one message before producer polling is paused.
///
/// The lower [`P2PStream`] sink admits two outbound messages and rejects uncompressed payloads
/// above 16 MiB, so 32 MiB mirrors the largest payload volume the lower p2p layer is already
/// prepared to buffer.
const MAX_MUX_OUT_BUFFER_BYTES: usize = 32 * 1024 * 1024;

#[derive(Debug)]
struct OutBuffer {
    messages: VecDeque<Bytes>,
    bytes: usize,
    max_bytes: usize,
}

impl Default for OutBuffer {
    fn default() -> Self {
        Self { messages: Default::default(), bytes: 0, max_bytes: MAX_MUX_OUT_BUFFER_BYTES }
    }
}

impl OutBuffer {
    fn push_back(&mut self, msg: Bytes) {
        self.bytes += msg.len();
        self.messages.push_back(msg);
    }

    fn pop_front(&mut self) -> Option<Bytes> {
        let msg = self.messages.pop_front()?;
        self.bytes -= msg.len();
        Some(msg)
    }

    fn is_empty(&self) -> bool {
        self.messages.is_empty()
    }

    const fn is_full(&self) -> bool {
        self.bytes >= self.max_bytes
    }
}

/// Result of polling outbound producers into the mux buffer.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum ProducerPoll {
    /// All outbound producers are pending.
    Pending,
    /// The mux buffer reached its soft cap.
    Full,
    /// An outbound producer closed.
    Closed,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        handshake::EthHandshake,
        message::MAX_MESSAGE_SIZE,
        protocol::Protocol,
        test_utils::{
            connect_passthrough, eth_handshake, eth_hello,
            proto::{test_hello, TestProtoMessage},
        },
        UnauthedEthStream, UnauthedP2PStream,
    };
    use futures::{stream, task::noop_waker_ref};
    use reth_eth_wire_types::EthNetworkPrimitives;
    use std::task::Poll;
    use tokio::{net::TcpListener, sync::oneshot};
    use tokio_util::codec::Decoder;

    #[derive(Debug)]
    struct PendingPrimary {
        _proxy: ProtocolProxy,
    }

    impl Stream for PendingPrimary {
        type Item = Result<(), P2PStreamError>;

        fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
            Poll::Pending
        }
    }

    #[derive(Debug)]
    struct StalledTransport;

    impl Stream for StalledTransport {
        type Item = io::Result<BytesMut>;

        fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
            Poll::Pending
        }
    }

    impl Sink<Bytes> for StalledTransport {
        type Error = io::Error;

        fn poll_ready(
            self: Pin<&mut Self>,
            _cx: &mut Context<'_>,
        ) -> Poll<Result<(), Self::Error>> {
            Poll::Pending
        }

        fn start_send(self: Pin<&mut Self>, _item: Bytes) -> Result<(), Self::Error> {
            Ok(())
        }

        fn poll_flush(
            self: Pin<&mut Self>,
            _cx: &mut Context<'_>,
        ) -> Poll<Result<(), Self::Error>> {
            Poll::Pending
        }

        fn poll_close(
            self: Pin<&mut Self>,
            _cx: &mut Context<'_>,
        ) -> Poll<Result<(), Self::Error>> {
            Poll::Pending
        }
    }

    #[tokio::test]
    async fn satellite_mux_stops_polling_protocols_when_out_buffer_is_full() {
        let (hello, _) = test_hello();
        let shared_capabilities =
            SharedCapabilities::try_new(hello.protocols.clone(), hello.message().capabilities)
                .unwrap();
        let conn = P2PStream::new(StalledTransport, shared_capabilities);
        let eth = conn.shared_capabilities().eth().unwrap().clone();

        let mut st = RlpxProtocolMultiplexer::new(conn)
            .into_satellite_stream(eth.capability().as_ref(), |proxy| PendingPrimary {
                _proxy: proxy,
            })
            .unwrap();
        const MESSAGE_COUNT: usize = 4096;
        const MESSAGE_BYTES: usize = 1024;
        st.inner.out_buffer.max_bytes = 4 * MESSAGE_BYTES + 1;
        st.install_protocol(&TestProtoMessage::capability(), |_conn| {
            stream::iter((0..MESSAGE_COUNT).map(|_| {
                let mut msg = BytesMut::zeroed(MESSAGE_BYTES);
                msg[0] = TestProtoMessage::ping().message_type as u8;
                msg
            }))
        })
        .unwrap();

        let mut cx = Context::from_waker(noop_waker_ref());
        assert!(Pin::new(&mut st).poll_next(&mut cx).is_pending());

        assert!(st.inner.out_buffer.bytes > st.inner.out_buffer.max_bytes);
        assert!(st.inner.out_buffer.bytes <= st.inner.out_buffer.max_bytes + MESSAGE_BYTES);
        assert!(st.inner.out_buffer.messages.len() < MESSAGE_COUNT);
    }

    #[tokio::test]
    async fn satellite_mux_round_robins_ready_protocols_when_out_buffer_fills() {
        let (mut hello, _) = eth_hello();
        let cap_a = Capability::new_static("aaa", 1);
        let cap_b = Capability::new_static("bbb", 1);
        hello.protocols.push(Protocol::new(cap_a.clone(), 1));
        hello.protocols.push(Protocol::new(cap_b.clone(), 1));

        let shared_capabilities =
            SharedCapabilities::try_new(hello.protocols.clone(), hello.message().capabilities)
                .unwrap();
        let conn = P2PStream::new(StalledTransport, shared_capabilities);
        let eth = conn.shared_capabilities().eth().unwrap().clone();
        let cap_a_offset =
            conn.shared_capabilities().find(&cap_a).unwrap().relative_message_id_offset();
        let cap_b_offset =
            conn.shared_capabilities().find(&cap_b).unwrap().relative_message_id_offset();

        let mut st = RlpxProtocolMultiplexer::new(conn)
            .into_satellite_stream(eth.capability().as_ref(), |proxy| PendingPrimary {
                _proxy: proxy,
            })
            .unwrap();
        st.inner.out_buffer.max_bytes = 5;
        st.install_protocol(&cap_a, |_conn| {
            stream::iter((0..16).map(|_| BytesMut::from(&[0, b'a'][..])))
        })
        .unwrap();
        st.install_protocol(&cap_b, |_conn| {
            stream::iter((0..16).map(|_| BytesMut::from(&[0, b'b'][..])))
        })
        .unwrap();

        let mut cx = Context::from_waker(noop_waker_ref());
        assert!(Pin::new(&mut st).poll_next(&mut cx).is_pending());

        let message_ids =
            st.inner.out_buffer.messages.iter().take(2).map(|msg| msg[0]).collect::<Vec<_>>();
        assert_eq!(message_ids.len(), 2);
        assert_ne!(message_ids[0], message_ids[1]);
        assert!(message_ids.contains(&cap_a_offset));
        assert!(message_ids.contains(&cap_b_offset));
    }

    #[tokio::test]
    async fn eth_satellite() {
        reth_tracing::init_test_tracing();
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let local_addr = listener.local_addr().unwrap();
        let (status, fork_filter) = eth_handshake();
        let other_status = status;
        let other_fork_filter = fork_filter.clone();
        let _handle = tokio::spawn(async move {
            let (incoming, _) = listener.accept().await.unwrap();
            let stream = crate::PassthroughCodec::default().framed(incoming);
            let (server_hello, _) = eth_hello();
            let (p2p_stream, _) =
                UnauthedP2PStream::new(stream).handshake(server_hello).await.unwrap();

            let (_eth_stream, _) = UnauthedEthStream::new(p2p_stream)
                .handshake::<EthNetworkPrimitives>(other_status, other_fork_filter)
                .await
                .unwrap();

            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
        });

        let conn = connect_passthrough(local_addr, eth_hello().0).await;
        let eth = conn.shared_capabilities().eth().unwrap().clone();

        let multiplexer = RlpxProtocolMultiplexer::new(conn);
        let _satellite = multiplexer
            .into_satellite_stream_with_handshake(eth.capability().as_ref(), async move |proxy| {
                UnauthedEthStream::new(proxy)
                    .handshake::<EthNetworkPrimitives>(status, fork_filter)
                    .await
            })
            .await
            .unwrap();
    }

    /// A test that install a satellite stream eth+test protocol and sends messages between them.
    #[tokio::test(flavor = "multi_thread")]
    async fn eth_test_protocol_satellite() {
        reth_tracing::init_test_tracing();
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let local_addr = listener.local_addr().unwrap();
        let (status, fork_filter) = eth_handshake();
        let other_status = status;
        let other_fork_filter = fork_filter.clone();
        let _handle = tokio::spawn(async move {
            let (incoming, _) = listener.accept().await.unwrap();
            let stream = crate::PassthroughCodec::default().framed(incoming);
            let (server_hello, _) = test_hello();
            let (conn, _) = UnauthedP2PStream::new(stream).handshake(server_hello).await.unwrap();

            let (mut st, _their_status) = RlpxProtocolMultiplexer::new(conn)
                .into_eth_satellite_stream::<EthNetworkPrimitives>(
                    other_status,
                    other_fork_filter,
                    Arc::new(EthHandshake::default()),
                    MAX_MESSAGE_SIZE,
                )
                .await
                .unwrap();

            st.install_protocol(&TestProtoMessage::capability(), |mut conn| {
                async_stream::stream! {
                    yield TestProtoMessage::ping().encoded();
                    let msg = conn.next().await.unwrap();
                    let msg = TestProtoMessage::decode_message(&mut &msg[..]).unwrap();
                    assert_eq!(msg, TestProtoMessage::pong());

                    yield TestProtoMessage::message("hello").encoded();
                    let msg = conn.next().await.unwrap();
                    let msg = TestProtoMessage::decode_message(&mut &msg[..]).unwrap();
                    assert_eq!(msg, TestProtoMessage::message("good bye!"));

                    yield TestProtoMessage::message("good bye!").encoded();

                    futures::future::pending::<()>().await;
                    unreachable!()
                }
            })
            .unwrap();

            loop {
                let _ = st.next().await;
            }
        });

        let conn = connect_passthrough(local_addr, test_hello().0).await;
        let (mut st, _their_status) = RlpxProtocolMultiplexer::new(conn)
            .into_eth_satellite_stream::<EthNetworkPrimitives>(
                status,
                fork_filter,
                Arc::new(EthHandshake::default()),
                MAX_MESSAGE_SIZE,
            )
            .await
            .unwrap();

        let (tx, mut rx) = oneshot::channel();

        st.install_protocol(&TestProtoMessage::capability(), |mut conn| {
            async_stream::stream! {
                let msg = conn.next().await.unwrap();
                let msg = TestProtoMessage::decode_message(&mut &msg[..]).unwrap();
                assert_eq!(msg, TestProtoMessage::ping());

                yield TestProtoMessage::pong().encoded();

                let msg = conn.next().await.unwrap();
                let msg = TestProtoMessage::decode_message(&mut &msg[..]).unwrap();
                assert_eq!(msg, TestProtoMessage::message("hello"));

                yield TestProtoMessage::message("good bye!").encoded();

                let msg = conn.next().await.unwrap();
                let msg = TestProtoMessage::decode_message(&mut &msg[..]).unwrap();
                assert_eq!(msg, TestProtoMessage::message("good bye!"));

                tx.send(()).unwrap();

                futures::future::pending::<()>().await;
                unreachable!()
            }
        })
        .unwrap();

        loop {
            tokio::select! {
                _ = &mut rx => {
                    break
                }
               _ = st.next() => {
                }
            }
        }
    }
}
