use crate::{
    Handshake, Upgrader,
    utils::codec::{Error as FrameError, recv_frame, validate_frame_len},
};
use commonware_codec::{
    Copying, DecodeExt, Encode, EncodeSize, FixedSize, Mode, Write, mode, modes, varint::UInt,
};
use commonware_cryptography::{Cipher, transcript::Transcript};
use commonware_runtime::{
    BufMut, Clock, Error as RuntimeError, IoBuf, IoBufMut, IoBufs, Sink, Stream,
};
use commonware_utils::Widen;
use rand_core::CryptoRng;
use std::{future::Future, marker::PhantomData};
use thiserror::Error;

/// Size of the length field in a version 1 header, excluding its tag.
const V1_HEADER_PLAINTEXT_SIZE: usize = u32::SIZE;

const CUPS_NAMESPACE: &[u8] = b"_COMMONWARE_STREAM_CUPS";
const LABEL_CIPHER_L2D: &[u8] = b"cipher_l2d";
const LABEL_CIPHER_D2L: &[u8] = b"cipher_d2l";

/// Returns the tag size of `C` as a record length.
const fn tag_size<C: Cipher>() -> u32 {
    let size = <C::Tag as FixedSize>::SIZE;
    assert!(size <= u32::MAX as usize);
    size as u32
}

/// Largest record payload that fits alongside its tag in a version 0 length prefix.
const fn max_size<C: Cipher>() -> u32 {
    u32::MAX - tag_size::<C>()
}

/// Returns the size of a version 1 header sealed by `C`.
const fn v1_header_size<C: Cipher>() -> usize {
    V1_HEADER_PLAINTEXT_SIZE + <C::Tag as FixedSize>::SIZE
}

/// Seals `data` with `cipher` and returns its tag, leaving the cipher empty if sealing fails.
fn seal<C: Cipher>(cipher: &mut Option<C>, data: &mut [u8]) -> Result<C::Tag, Error> {
    let sealer = cipher.take().ok_or(Error::StreamClosed)?;
    let (sealer, tag) = sealer.seal(&[], data).ok_or(Error::SealFailed)?;
    *cipher = Some(sealer);
    Ok(tag)
}

/// Opens `buf`, a ciphertext followed by its tag, with `cipher` and returns its plaintext, leaving
/// the cipher empty if opening fails.
fn open<C: Cipher>(cipher: &mut Option<C>, mut buf: IoBufMut) -> Result<IoBufMut, Error> {
    let opener = cipher.take().ok_or(Error::StreamClosed)?;
    let len = buf
        .len()
        .checked_sub(<C::Tag as FixedSize>::SIZE)
        .ok_or(Error::OpenFailed)?;
    let (data, tag) = buf.as_mut().split_at_mut(len);
    let tag = C::Tag::decode(Copying(&*tag)).map_err(|_| Error::OpenFailed)?;
    *cipher = Some(opener.open(&[], data, &tag).ok_or(Error::OpenFailed)?);
    buf.truncate(len);
    Ok(buf)
}

/// Errors that can occur when sending or receiving records.
#[derive(Error, Debug)]
pub enum Error {
    #[error("seal failed")]
    SealFailed,
    #[error("open failed")]
    OpenFailed,
    #[error("recv failed")]
    RecvFailed(RuntimeError),
    #[error("recv too large: {0} bytes")]
    RecvTooLarge(usize),
    #[error("invalid varint length prefix")]
    InvalidVarint,
    #[error("send failed")]
    SendFailed(RuntimeError),
    #[error("send too large: {0} bytes")]
    SendTooLarge(usize),
    #[error("closed after an earlier failure")]
    StreamClosed,
}

impl From<FrameError> for Error {
    fn from(value: FrameError) -> Self {
        match value {
            FrameError::RecvFailed(err) => Self::RecvFailed(err),
            FrameError::RecvTooLarge(len) => Self::RecvTooLarge(len),
            FrameError::InvalidVarint => Self::InvalidVarint,
            FrameError::SendFailed(err) => Self::SendFailed(err),
            FrameError::SendTooLarge(len) => Self::SendTooLarge(len),
        }
    }
}

/// Key derivation and record format used by [Cups].
///
/// Both peers must use the same version.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
pub enum Version {
    /// Records framed by a visible length prefix.
    V0,
    /// Records framed by an encrypted, authenticated length header.
    V1,
}

impl From<Version> for Mode {
    fn from(version: Version) -> Self {
        match version {
            Version::V0 => mode!(0),
            Version::V1 => mode!(1),
        }
    }
}

/// [Upgrader] that runs handshake `H`, then exchanges records of one [Version] sealed by `C`.
pub struct Cups<H, C> {
    /// Handshake used to authenticate peers and agree on a secret transcript.
    pub handshake: H,

    /// Version used for key derivation and record framing.
    version: Version,

    // `fn() -> C` keeps `Cups` `Send` and `Sync` for any `C`.
    _cipher: PhantomData<fn() -> C>,
}

// Manual implementation avoids the `C: Clone` bound a derive would add.
impl<H: Clone, C> Clone for Cups<H, C> {
    fn clone(&self) -> Self {
        Self {
            handshake: self.handshake.clone(),
            version: self.version,
            _cipher: PhantomData,
        }
    }
}

impl<H, C> Cups<H, C> {
    /// Runs `handshake`, then exchanges `version` records sealed by `C`.
    pub const fn new(handshake: H, version: Version) -> Self {
        Self {
            handshake,
            version,
            _cipher: PhantomData,
        }
    }
}

impl<H, C: Cipher> Cups<H, C> {
    /// Derives listener-to-dialer and dialer-to-listener ciphers from the secret transcript.
    fn ciphers(mut transcript: Transcript, version: Version) -> (C, C) {
        if let Some(modes) = modes![version] {
            transcript = transcript.fork(CUPS_NAMESPACE);
            transcript.commit(modes.encode());
        }
        (
            C::random(transcript.noise(LABEL_CIPHER_L2D)),
            C::random(transcript.noise(LABEL_CIPHER_D2L)),
        )
    }

    /// Largest supported plaintext message, in bytes.
    ///
    /// Version 0 prefixes count the tag, so the payload and tag must fit in a `u32`.
    pub const MAX_SIZE: u32 = max_size::<C>();

    /// Returns the size of the header that precedes a record carrying `len` payload bytes.
    ///
    /// # Panics
    ///
    /// Panics if `len` exceeds the largest record payload, `u32::MAX` minus the tag size.
    pub fn header_len(&self, len: usize) -> usize {
        let len = u32::try_from(len)
            .ok()
            .filter(|len| *len <= Self::MAX_SIZE)
            .expect("payload exceeds stream limit");
        match self.version {
            Version::V0 => UInt(len + tag_size::<C>()).encode_size(),
            Version::V1 => v1_header_size::<C>(),
        }
    }

    /// Returns the encoded size of a record carrying `len` payload bytes.
    ///
    /// # Panics
    ///
    /// Panics if `len` exceeds the largest record payload, `u32::MAX` minus the tag size, or if the
    /// record size does not fit in a `usize`.
    pub fn record_len(&self, len: usize) -> usize {
        self.header_len(len)
            .checked_add(len)
            .and_then(|size| size.checked_add(<C::Tag as FixedSize>::SIZE))
            .expect("record size exceeds usize")
    }

    /// Splits the derived ciphers into the sender and receiver halves of a connection.
    const fn split<I: Stream, O: Sink>(
        version: Version,
        send: C,
        recv: C,
        stream: I,
        sink: O,
        max_message_size: u32,
    ) -> (Sender<C, O>, Receiver<C, I>) {
        (
            Sender {
                cipher: Some(send),
                sink,
                max_message_size,
                version,
            },
            Receiver {
                cipher: Some(recv),
                stream,
                max_message_size,
                version,
            },
        )
    }
}

impl<H: Handshake, C: Cipher> Upgrader for Cups<H, C> {
    const MAX_SIZE: u32 = max_size::<C>();

    type PublicKey = H::PublicKey;
    type Error = H::Error;
    type Sender<I: Stream, O: Sink> = Sender<C, O>;
    type Receiver<I: Stream, O: Sink> = Receiver<C, I>;

    fn public_key(&self) -> Self::PublicKey {
        self.handshake.public_key()
    }

    async fn dial<E, I, O>(
        self,
        context: E,
        namespace: &[u8],
        max_message_size: u32,
        peer: Self::PublicKey,
        mut stream: I,
        mut sink: O,
    ) -> Result<(Self::Sender<I, O>, Self::Receiver<I, O>), Self::Error>
    where
        E: Clock + CryptoRng,
        I: Stream,
        O: Sink,
    {
        assert!(
            max_message_size <= Self::MAX_SIZE,
            "maximum message size exceeds stream limit"
        );
        let version = self.version;
        let transcript = self
            .handshake
            .dial(context, namespace, peer, &mut stream, &mut sink)
            .await?;
        let (recv, send) = Self::ciphers(transcript, version);
        Ok(Self::split(
            version,
            send,
            recv,
            stream,
            sink,
            max_message_size,
        ))
    }

    async fn listen<E, I, O, B, F>(
        self,
        context: E,
        namespace: &[u8],
        max_message_size: u32,
        bouncer: B,
        mut stream: I,
        mut sink: O,
    ) -> Result<(Self::PublicKey, Self::Sender<I, O>, Self::Receiver<I, O>), Self::Error>
    where
        E: Clock + CryptoRng,
        I: Stream,
        O: Sink,
        B: FnOnce(Self::PublicKey) -> F + Send,
        F: Future<Output = bool> + Send,
    {
        assert!(
            max_message_size <= Self::MAX_SIZE,
            "maximum message size exceeds stream limit"
        );
        let version = self.version;
        let (peer, transcript) = self
            .handshake
            .listen(context, namespace, bouncer, &mut stream, &mut sink)
            .await?;
        let (send, recv) = Self::ciphers(transcript, version);
        let (sender, receiver) = Self::split(version, send, recv, stream, sink, max_message_size);
        Ok((peer, sender, receiver))
    }
}

/// Appends the header for an encrypted payload of `len` bytes. A V1 header is sealed and consumes
/// a cipher position.
fn append_header<C: Cipher>(
    version: Version,
    buf: &mut IoBufMut,
    cipher: &mut Option<C>,
    len: u32,
) -> Result<(), Error> {
    match version {
        // Callers bound `len` by `MAX_SIZE`, so adding the tag cannot overflow.
        Version::V0 => UInt(len + tag_size::<C>()).write(buf),
        Version::V1 => {
            let offset = buf.len();
            len.write(buf);
            let tag = seal(cipher, &mut buf.as_mut()[offset..])?;
            buf.put_slice(&tag);
        }
    }
    Ok(())
}

/// Receives the next record's encrypted payload and tag after validating its header.
async fn recv_record<C: Cipher>(
    version: Version,
    stream: &mut impl Stream,
    cipher: &mut Option<C>,
    max_message_size: u32,
) -> Result<IoBufs, Error> {
    match version {
        Version::V0 => recv_frame(stream, max_message_size.saturating_add(tag_size::<C>()))
            .await
            .map_err(|err| match err {
                // The prefix counts the tag, which is excluded from the reported payload.
                FrameError::RecvTooLarge(len) => {
                    Error::RecvTooLarge(len - <C::Tag as FixedSize>::SIZE)
                }
                err => err.into(),
            }),
        Version::V1 => {
            // Request the fixed-size header before trusting the payload length. The header is
            // decrypted in its received buffer when possible.
            let header = stream
                .recv(v1_header_size::<C>())
                .await
                .map_err(Error::RecvFailed)?;

            // Authenticate the header before decoding its length or requesting the payload.
            let header = open(cipher, mutable_frame(header))?;
            assert_eq!(header.len(), V1_HEADER_PLAINTEXT_SIZE);
            let len = u32::decode(header).expect("header has a fixed size");
            if len > max_message_size {
                return Err(Error::RecvTooLarge(Widen::<usize>::widen(len)));
            }

            // Receive the payload and tag only after authenticating and checking the length.
            let body_len = Widen::<usize>::widen(len) + <C::Tag as FixedSize>::SIZE;
            stream.recv(body_len).await.map_err(Error::RecvFailed)
        }
    }
}

/// Sends CUPS records to a peer.
pub struct Sender<C, O> {
    cipher: Option<C>,
    sink: O,
    max_message_size: u32,
    version: Version,
}

impl<C: Cipher, O: Sink> Sender<C, O> {
    /// Returns the encoded size of a frame carrying `len` payload bytes, or
    /// [Error::SendTooLarge] if `len` exceeds the message limit.
    ///
    /// The returned size includes the header, ciphertext, and AEAD tags.
    fn encrypted_frame_len(&self, len: usize) -> Result<usize, Error> {
        validate_frame_len(len, self.max_message_size)?;
        Ok(Cups::<(), C>::new((), self.version).record_len(len))
    }

    /// Appends one encrypted frame carrying `bufs` to `buf`.
    ///
    /// Appending in place lets [`Self::encrypt_frames`] pack consecutive frames into one allocation
    /// without staging each frame in its own buffer.
    fn append_encrypted_frame(
        &mut self,
        buf: &mut IoBufMut,
        mut bufs: IoBufs,
    ) -> Result<(), Error> {
        let len = validate_frame_len(bufs.len(), self.max_message_size)?;
        append_header(self.version, buf, &mut self.cipher, len)?;

        // Copy the plaintext directly into the frame.
        let plaintext_offset = buf.len();
        buf.put(&mut bufs);

        // Encrypt in-place and append the tag to the frame.
        let tag = seal(&mut self.cipher, &mut buf.as_mut()[plaintext_offset..])?;
        buf.put_slice(&tag);
        assert_eq!(
            buf.len() - plaintext_offset,
            Widen::<usize>::widen(len) + <C::Tag as FixedSize>::SIZE
        );
        Ok(())
    }

    /// Encrypts each message as one frame and returns the frames in one contiguous buffer.
    ///
    /// `total_len` must be the combined encoded size of the frames, so the buffer is allocated
    /// once at exactly that size.
    fn encrypt_frames<I>(&mut self, messages: I, total_len: usize) -> Result<IoBuf, Error>
    where
        I: IntoIterator<Item = IoBufs>,
    {
        let mut frames = IoBufMut::with_capacity(total_len);
        for msg in messages {
            self.append_encrypted_frame(&mut frames, msg)?;
        }
        assert_eq!(frames.len(), total_len);
        Ok(frames.freeze())
    }

    /// Encrypts and sends a message to the peer.
    ///
    /// Copies the message into a new buffer, encrypts it in place, and sends the record in one sink
    /// call.
    pub async fn send(&mut self, bufs: impl Into<IoBufs>) -> Result<(), Error> {
        let bufs = bufs.into();
        let frame_len = self.encrypted_frame_len(bufs.len())?;
        let frames = self.encrypt_frames(std::iter::once(bufs), frame_len)?;
        self.sink.send(frames).await.map_err(Error::SendFailed)
    }

    /// Encrypts and sends messages in a single sink call, one record per message.
    ///
    /// If any message exceeds the message limit, the batch fails with [Error::SendTooLarge] before
    /// any record is sealed. No cipher position is consumed, so the sender remains usable. An empty
    /// batch succeeds without calling the sink.
    pub async fn send_many<B, I>(&mut self, bufs: I) -> Result<(), Error>
    where
        B: Into<IoBufs>,
        I: IntoIterator<Item = B>,
    {
        // A saturating total cannot overflow, so the sizing pass reaches every record and reports
        // any oversized one before anything is sealed.
        let bufs = bufs.into_iter();
        let mut messages = Vec::with_capacity(bufs.size_hint().0);
        let mut total_len = 0usize;
        for buf in bufs {
            let msg = buf.into();
            total_len = total_len.saturating_add(self.encrypted_frame_len(msg.len())?);
            messages.push(msg);
        }
        if messages.is_empty() {
            return Ok(());
        }

        // Seal the records back to back into one buffer, so the sink receives a single chunk.
        let frames = self.encrypt_frames(messages, total_len)?;
        self.sink.send(frames).await.map_err(Error::SendFailed)
    }
}

/// Receives CUPS records from a peer.
pub struct Receiver<C, I> {
    cipher: Option<C>,
    stream: I,
    max_message_size: u32,
    version: Version,
}

impl<C: Cipher, O: Sink> crate::Sender for Sender<C, O> {
    type Error = Error;

    async fn send(&mut self, message: impl Into<IoBufs> + Send) -> Result<(), Self::Error> {
        Self::send(self, message).await
    }

    async fn send_many<I>(&mut self, messages: I) -> Result<(), Self::Error>
    where
        I: IntoIterator + Send,
        I::Item: Into<IoBufs> + Send,
        I::IntoIter: Send,
    {
        Self::send_many(self, messages).await
    }
}

impl<C: Cipher, I: Stream> crate::Receiver for Receiver<C, I> {
    type Error = Error;

    async fn recv(&mut self) -> Result<IoBufs, Self::Error> {
        Self::recv(self).await
    }
}

/// Returns a received frame as one mutable buffer for in-place decryption. The frame is copied
/// only when it spans several buffers or [`IoBuf::try_into_mut`] cannot recover its one buffer.
fn mutable_frame(encrypted: IoBufs) -> IoBufMut {
    match encrypted
        .try_into_single()
        .and_then(|buf| buf.try_into_mut().map_err(IoBufs::from))
    {
        Ok(buf) => buf,
        Err(mut encrypted) => {
            let mut buf = IoBufMut::with_capacity(encrypted.len());
            buf.put(&mut encrypted);
            buf
        }
    }
}

impl<C: Cipher, I: Stream> Receiver<C, I> {
    /// Receives and decrypts a message from the peer.
    ///
    /// Receives ciphertext and decrypts it in place when the received frame is a single buffer that
    /// [`IoBuf::try_into_mut`] can recover. Otherwise, allocates a buffer, copies the ciphertext,
    /// and decrypts the copy in place.
    ///
    /// Any error closes the receiver, so later calls return [Error::StreamClosed] without reading
    /// from the stream.
    pub async fn recv(&mut self) -> Result<IoBufs, Error> {
        // An earlier error closed this half, so refuse without reading from the stream.
        if self.cipher.is_none() {
            return Err(Error::StreamClosed);
        }

        // Every receive error is terminal for the connection, and some leave the cipher position
        // ahead of the stream, so refuse later records instead of reading them.
        let result = self.recv_open().await;
        if result.is_err() {
            self.cipher = None;
        }
        result
    }

    /// Receives the next record and decrypts it.
    async fn recv_open(&mut self) -> Result<IoBufs, Error> {
        let encrypted = recv_record(
            self.version,
            &mut self.stream,
            &mut self.cipher,
            self.max_message_size,
        )
        .await?;

        // Decrypt in place, keeping only the plaintext.
        let plaintext = open(&mut self.cipher, mutable_frame(encrypted))?;
        Ok(plaintext.freeze().into())
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use commonware_cryptography::{ChaCha20Poly1305, transcript::Version as TranscriptVersion};
    use commonware_math::algebra::Random;
    use commonware_runtime::{Runner as _, Supervisor as _, deterministic, mocks};
    use commonware_utils::TestRng;
    use futures::FutureExt as _;

    /// Message limit, in bytes, of the test senders and receivers.
    const MAX_MESSAGE_SIZE: u32 = 64 * 1024;

    type RecordCipher = ChaCha20Poly1305;
    type Tag = <RecordCipher as Cipher>::Tag;
    type TestCups = Cups<(), RecordCipher>;
    const TAG_SIZE: u32 = tag_size::<RecordCipher>();
    const MAX_SIZE: u32 = TestCups::MAX_SIZE;

    /// Supplies a deterministic test transcript to both roles, or rejects the handshake.
    #[derive(Clone)]
    struct TestHandshake(Option<TranscriptVersion>);

    impl TestHandshake {
        fn transcript(&self, namespace: &[u8]) -> Result<Transcript, std::io::Error> {
            let framing = self.0.ok_or_else(|| std::io::Error::other("rejected"))?;
            let mut transcript = Transcript::new(namespace, framing);
            transcript.commit(b"test shared secret".as_slice());
            Ok(transcript)
        }
    }

    impl Handshake for TestHandshake {
        type PublicKey = ();
        type Error = std::io::Error;

        fn public_key(&self) -> Self::PublicKey {}

        async fn dial<E, I, O>(
            self,
            _context: E,
            namespace: &[u8],
            _peer: (),
            _stream: &mut I,
            _sink: &mut O,
        ) -> Result<Transcript, Self::Error>
        where
            E: Clock + CryptoRng,
            I: Stream,
            O: Sink,
        {
            assert_eq!(namespace, b"application");
            self.transcript(namespace)
        }

        async fn listen<E, I, O, B, F>(
            self,
            _context: E,
            namespace: &[u8],
            bouncer: B,
            _stream: &mut I,
            _sink: &mut O,
        ) -> Result<(Self::PublicKey, Transcript), Self::Error>
        where
            E: Clock + CryptoRng,
            I: Stream,
            O: Sink,
            B: FnOnce(Self::PublicKey) -> F + Send,
            F: Future<Output = bool> + Send,
        {
            assert_eq!(namespace, b"application");
            if !bouncer(()).await {
                return Err(std::io::Error::other("rejected"));
            }
            Ok(((), self.transcript(namespace)?))
        }
    }

    /// Pins the CUPS fork, canonical mode, and directional cipher derivation under both framings.
    #[test]
    fn test_cipher_derivation() {
        for framing in [TranscriptVersion::V0, TranscriptVersion::V1] {
            for version in [Version::V0, Version::V1] {
                // Derive the traffic ciphers through the CUPS setup path.
                let transcript = TestHandshake(Some(framing))
                    .transcript(b"application")
                    .unwrap();
                let (actual_l2d, actual_d2l) = TestCups::ciphers(transcript, version);

                // Build an oracle from the specified namespace, mode, and direction labels.
                let mut expected = TestHandshake(Some(framing))
                    .transcript(b"application")
                    .unwrap();
                if version == Version::V1 {
                    expected = expected.fork(b"_COMMONWARE_STREAM_CUPS");
                    expected.commit(&[1][..]);
                }
                let expected_l2d = RecordCipher::random(expected.noise(b"cipher_l2d"));
                let expected_d2l = RecordCipher::random(expected.noise(b"cipher_d2l"));

                // Each directional key must produce the oracle's ciphertext and authentication tag.
                for (actual, expected) in [(actual_l2d, expected_l2d), (actual_d2l, expected_d2l)] {
                    let mut actual = Some(actual);
                    let mut expected = Some(expected);
                    let mut a = b"oracle".to_vec();
                    let mut b = a.clone();
                    assert_eq!(
                        seal(&mut actual, &mut a).unwrap(),
                        seal(&mut expected, &mut b).unwrap()
                    );
                    assert_eq!(a, b);
                }
            }
        }
    }

    /// Checks that CUPS uses a custom handshake's transcript to exchange records in both directions.
    #[test]
    fn test_custom_handshake_round_trip() {
        for framing in [TranscriptVersion::V0, TranscriptVersion::V1] {
            for version in [Version::V0, Version::V1] {
                deterministic::Runner::default().start(|context| async move {
                    // Establish both CUPS roles from the same handshake transcript.
                    let (dial_sink, listen_stream) = mocks::Channel::init();
                    let (listen_sink, dial_stream) = mocks::Channel::init();
                    let cups = Cups::<_, RecordCipher>::new(TestHandshake(Some(framing)), version);
                    let (mut dial_sender, mut dial_receiver) = cups
                        .clone()
                        .dial(
                            context.child("dial"),
                            b"application",
                            64,
                            (),
                            dial_stream,
                            dial_sink,
                        )
                        .await
                        .unwrap();
                    let (_, mut listen_sender, mut listen_receiver) = cups
                        .listen(
                            context.child("listen"),
                            b"application",
                            64,
                            |_| async { true },
                            listen_stream,
                            listen_sink,
                        )
                        .await
                        .unwrap();

                    // Exchange records in both directions to check the send/receive key assignment.
                    dial_sender.send(b"dial".as_slice()).await.unwrap();
                    assert_eq!(listen_receiver.recv().await.unwrap().coalesce(), b"dial");
                    listen_sender.send(b"listen".as_slice()).await.unwrap();
                    assert_eq!(dial_receiver.recv().await.unwrap().coalesce(), b"listen");
                });
            }
        }
    }

    /// Checks that CUPS forwards the application namespace and handshake errors in either role.
    #[test]
    fn test_custom_handshake() {
        deterministic::Runner::default().start(|context| async move {
            for version in [Version::V0, Version::V1] {
                let cups = Cups::<_, RecordCipher>::new(TestHandshake(None), version);

                // Dialing runs the handshake with the application namespace and returns its error.
                let (sink, stream) = mocks::Channel::init();
                assert!(
                    cups.clone()
                        .dial(context.child("dial"), b"application", 1, (), stream, sink)
                        .await
                        .is_err()
                );

                // Listening does the same.
                let (sink, stream) = mocks::Channel::init();
                assert!(
                    cups.listen(
                        context.child("listen"),
                        b"application",
                        1,
                        |_| async { true },
                        stream,
                        sink
                    )
                    .await
                    .is_err()
                );
            }
        });
    }

    /// Returns the cipher every test peer derives, so each side replays the other's positions.
    fn cipher() -> Option<RecordCipher> {
        Some(RecordCipher::random(TestRng::new(0)))
    }

    /// Returns a sender of `version` records into `sink`.
    fn sender(
        sink: mocks::Sink,
        max_message_size: u32,
        version: Version,
    ) -> Sender<RecordCipher, mocks::Sink> {
        Sender {
            cipher: cipher(),
            sink,
            max_message_size,
            version,
        }
    }

    /// Returns a receiver of `version` records from `stream`.
    fn receiver(
        stream: mocks::Stream,
        max_message_size: u32,
        version: Version,
    ) -> Receiver<RecordCipher, mocks::Stream> {
        Receiver {
            cipher: cipher(),
            stream,
            max_message_size,
            version,
        }
    }

    /// Seals `msg` with `cipher` into a new buffer.
    fn sealed(cipher: &mut Option<ChaCha20Poly1305>, msg: &[u8]) -> Vec<u8> {
        let mut buf = msg.to_vec();
        let tag = seal(cipher, &mut buf).unwrap();
        buf.extend_from_slice(&tag);
        buf
    }

    /// Checks that each version writes a record as its length header followed by the sealed
    /// payload, and that `record_len` matches the bytes written.
    #[test]
    fn test_record_encoding() {
        for version in [Version::V0, Version::V1] {
            deterministic::Runner::default().start(|_| async move {
                let (sink, mut stream) = mocks::Channel::init();
                let mut sender = sender(sink, 64, version);
                let messages = [&b""[..], &b"hello"[..], &[7; 64][..]];
                sender.send_many(messages).await.unwrap();

                // A cipher from the same seed reproduces the sender's key and nonce sequence.
                let mut cipher = cipher();
                for message in messages {
                    let len = message.len() as u32;

                    // V0 prefixes a plaintext length that counts the tag. V1 prefixes a sealed
                    // length that excludes the tag.
                    let mut expected = match version {
                        Version::V0 => UInt(len + TAG_SIZE).encode().to_vec(),
                        Version::V1 => sealed(&mut cipher, &len.to_be_bytes()),
                    };
                    expected.extend(sealed(&mut cipher, message));
                    assert_eq!(
                        TestCups::new((), version).record_len(message.len()),
                        expected.len()
                    );
                    assert_eq!(
                        stream.recv(expected.len()).await.unwrap().coalesce(),
                        expected.as_slice()
                    );
                }
            });
        }
    }

    /// Checks that the largest payload leaves room for the tag in a `u32` record length.
    #[test]
    fn test_max_size_fits_tag() {
        assert_eq!(MAX_SIZE + TAG_SIZE, u32::MAX);
    }

    /// Checks that sizing a payload beyond the stream limit panics.
    #[test]
    #[should_panic(expected = "payload exceeds stream limit")]
    fn test_header_len_rejects_oversized_payload() {
        TestCups::new((), Version::V1).header_len(Widen::<usize>::widen(u32::MAX));
    }

    /// Checks that a version 0 receiver reports the payload length of an oversized record, which
    /// excludes the tag that its length prefix counts.
    #[test]
    fn test_v0_recv_too_large_reports_payload_len() {
        deterministic::Runner::default().start(|_| async move {
            let (mut sink, stream) = mocks::Channel::init();
            let mut receiver = receiver(stream, MAX_MESSAGE_SIZE, Version::V0);
            let len = MAX_MESSAGE_SIZE + 1;
            let prefix = UInt(len + TAG_SIZE).encode().to_vec();
            sink.send(prefix).await.unwrap();

            // The sink stays open without sending a payload, so rejection must not wait for it.
            let result = receiver
                .recv()
                .now_or_never()
                .expect("prefix rejection must be immediate");
            assert!(matches!(
                result,
                Err(Error::RecvTooLarge(n)) if n == Widen::<usize>::widen(len)
            ));
        });
    }

    /// Checks that frame failures surface as the matching stream errors.
    #[test]
    fn test_frame_errors_surface_as_stream_errors() {
        deterministic::Runner::default().start(|_| async move {
            // A length prefix that never terminates is an invalid varint.
            let (mut sink, stream) = mocks::Channel::init();
            let mut invalid = receiver(stream, MAX_MESSAGE_SIZE, Version::V0);
            sink.send(vec![0xFF; 5]).await.unwrap();
            assert!(matches!(invalid.recv().await, Err(Error::InvalidVarint)));

            // A closed peer fails the read of the length prefix.
            let (sink, stream) = mocks::Channel::init();
            let mut closed = receiver(stream, MAX_MESSAGE_SIZE, Version::V0);
            drop(sink);
            assert!(matches!(closed.recv().await, Err(Error::RecvFailed(_))));
        });
    }

    /// Checks that a version 0 record shorter than a tag fails to open.
    #[test]
    fn test_v0_record_shorter_than_tag_rejected() {
        deterministic::Runner::default().start(|_| async move {
            let (mut sink, stream) = mocks::Channel::init();
            let mut receiver = receiver(stream, MAX_MESSAGE_SIZE, Version::V0);
            let mut record = UInt(TAG_SIZE - 1).encode().to_vec();
            record.resize(record.len() + Tag::SIZE - 1, 0);
            sink.send(record).await.unwrap();
            assert!(matches!(receiver.recv().await, Err(Error::OpenFailed)));
        });
    }

    /// Checks that every receive error closes the receiver, so the next call returns
    /// `StreamClosed` without reading from the stream.
    #[test]
    fn test_recv_error_closes_receiver() {
        #[derive(Clone, Copy, Debug)]
        enum Expected {
            OpenFailed,
            RecvTooLarge,
            InvalidVarint,
            RecvFailed,
        }

        // Each case sends one malformed record: a corrupted empty record, an oversized length, or
        // a prefix that is not a varint. A case without bytes closes the sink instead.
        let oversized = MAX_MESSAGE_SIZE + 1;
        let v0_corrupt = {
            let mut record = UInt(TAG_SIZE).encode().to_vec();
            record.extend(sealed(&mut cipher(), b""));
            record[1] ^= 1;
            record
        };
        let v1_corrupt = {
            let mut header = sealed(&mut cipher(), &0u32.to_be_bytes());
            header[0] ^= 1;
            header
        };
        let cases = [
            (Version::V0, Some(v0_corrupt), Expected::OpenFailed),
            (
                Version::V0,
                Some(UInt(oversized + TAG_SIZE).encode().to_vec()),
                Expected::RecvTooLarge,
            ),
            (Version::V0, Some(vec![0xFF; 5]), Expected::InvalidVarint),
            (Version::V0, None, Expected::RecvFailed),
            (Version::V1, Some(v1_corrupt), Expected::OpenFailed),
            (
                Version::V1,
                Some(sealed(&mut cipher(), &oversized.to_be_bytes())),
                Expected::RecvTooLarge,
            ),
            (Version::V1, None, Expected::RecvFailed),
        ];
        for (version, bytes, expected) in cases {
            deterministic::Runner::default().start(|_| async move {
                let (mut sink, stream) = mocks::Channel::init();
                let mut receiver = receiver(stream, MAX_MESSAGE_SIZE, version);
                let _open_sink = match bytes {
                    Some(bytes) => {
                        sink.send(bytes).await.unwrap();
                        Some(sink)
                    }
                    None => {
                        drop(sink);
                        None
                    }
                };

                let result = receiver.recv().await;
                assert!(
                    match expected {
                        Expected::OpenFailed => matches!(result, Err(Error::OpenFailed)),
                        Expected::RecvTooLarge => {
                            matches!(result, Err(Error::RecvTooLarge(n)) if n == oversized as usize)
                        }
                        Expected::InvalidVarint => matches!(result, Err(Error::InvalidVarint)),
                        Expected::RecvFailed => matches!(result, Err(Error::RecvFailed(_))),
                    },
                    "{version:?}/{expected:?}: {:?}",
                    result.as_ref().err()
                );

                // Nothing more arrives. A receiver that still read from the stream would block on
                // the open sink or report its closure instead of refusing.
                assert!(
                    matches!(
                        receiver.recv().now_or_never(),
                        Some(Err(Error::StreamClosed))
                    ),
                    "{version:?}/{expected:?}"
                );
            });
        }
    }

    /// Checks that a sender without a cipher refuses every record without writing to the sink.
    #[test]
    fn test_send_after_failure_closed() {
        for version in [Version::V0, Version::V1] {
            deterministic::Runner::default().start(|_| async move {
                let (sink, mut stream) = mocks::Channel::init();

                // A failed seal leaves the sender without a cipher.
                let mut sender: Sender<RecordCipher, _> = Sender {
                    cipher: None,
                    sink,
                    max_message_size: MAX_MESSAGE_SIZE,
                    version,
                };

                // Both send paths refuse the record.
                assert!(matches!(
                    sender.send(&b"x"[..]).await,
                    Err(Error::StreamClosed)
                ));
                assert!(matches!(
                    sender.send_many([&b"x"[..]]).await,
                    Err(Error::StreamClosed)
                ));

                // Nothing reached the sink.
                assert!(stream.recv(1).now_or_never().is_none());
            });
        }
    }

    /// Checks that a sender reports the payload length of an oversized record without consuming
    /// its cipher or writing to the sink.
    #[test]
    fn test_send_too_large_reports_payload_len() {
        for version in [Version::V0, Version::V1] {
            deterministic::Runner::default().start(|_| async move {
                let (sink, mut stream) = mocks::Channel::init();
                let mut sender = sender(sink, MAX_MESSAGE_SIZE, version);
                let oversized = vec![0u8; MAX_MESSAGE_SIZE as usize + 1];

                // Both send paths report the payload length, excluding the tag.
                assert!(matches!(
                    sender.send(oversized.clone()).await,
                    Err(Error::SendTooLarge(n)) if n == oversized.len()
                ));
                assert!(matches!(
                    sender.send_many([oversized.clone()]).await,
                    Err(Error::SendTooLarge(n)) if n == oversized.len()
                ));

                // Nothing reached the sink, and the cipher is still usable.
                assert!(stream.recv(1).now_or_never().is_none());
                assert!(sender.cipher.is_some());
            });
        }
    }

    /// Checks that `send_many` sizes the whole batch before allocating its output buffer or sealing
    /// any record. A trailing oversized record is reported even when the frames before it total
    /// more than 4 GiB.
    #[test]
    fn test_send_many_sizes_whole_batch_before_allocating() {
        for version in [Version::V0, Version::V1] {
            deterministic::Runner::default().start(|_| async move {
                let (sink, mut stream) = mocks::Channel::init();
                let limit = 65_536;
                let mut sender = sender(sink, limit, version);

                // Repeat one valid message 65,536 times through shared views, so the frames total
                // more than 4 GiB and exceed any 32-bit length. One message over the limit ends
                // the batch.
                let valid = IoBufs::from(IoBuf::from(vec![0u8; limit as usize]));
                let oversized = IoBufs::from(IoBuf::from(vec![0u8; limit as usize + 1]));
                let batch = std::iter::repeat_n(valid, 65_536).chain(std::iter::once(oversized));
                assert!(matches!(
                    sender.send_many(batch).await,
                    Err(Error::SendTooLarge(n)) if n == limit as usize + 1
                ));

                // Nothing reached the sink, and the cipher is still usable.
                assert!(stream.recv(1).now_or_never().is_none());
                assert!(sender.cipher.is_some());
            });
        }
    }

    /// Checks that a version 1 receiver rejects an invalid header before any payload arrives.
    #[test]
    fn test_invalid_header_rejected_before_body() {
        // Cases: a flipped length byte, a flipped tag byte, a length one above the limit, and
        // `u32::MAX`.
        for (length, corrupt) in [
            (0, Some(0)),
            (0, Some(V1_HEADER_PLAINTEXT_SIZE)),
            (MAX_MESSAGE_SIZE + 1, None),
            (u32::MAX, None),
        ] {
            deterministic::Runner::default().start(|_| async move {
                let (mut sink, stream) = mocks::Channel::init();
                let mut receiver = receiver(stream, MAX_MESSAGE_SIZE, Version::V1);
                let mut cipher = cipher();
                let mut header = sealed(&mut cipher, &length.to_be_bytes());
                if let Some(offset) = corrupt {
                    header[offset] ^= 1;
                }
                sink.send(header).await.unwrap();

                // The sink stays open without sending a payload, so rejection must not wait for it.
                let result = receiver
                    .recv()
                    .now_or_never()
                    .expect("header rejection must be immediate");
                if corrupt.is_some() {
                    assert!(matches!(result, Err(Error::OpenFailed)));
                } else {
                    assert!(matches!(
                        result,
                        Err(Error::RecvTooLarge(n)) if n == Widen::<usize>::widen(length)
                    ));
                }
            });
        }
    }
}