unb-runtime 2.0.3

use std::future::poll_fn;
use std::pin::Pin;

use bytes::{Bytes, BytesMut};
use futures_util::Stream;
use unb_core::CoreError;

pub type BodyStream = Pin<Box<dyn Stream<Item = Result<Bytes, CoreError>> + Send>>;

pub enum WireBody {
    Bytes(Bytes),
    Stream(BodyStream),
}

impl WireBody {
    pub async fn next_chunk(&mut self) -> Option<Result<Bytes, CoreError>> {
        match self {
            WireBody::Bytes(bytes) if bytes.is_empty() => None,
            WireBody::Bytes(bytes) => Some(Ok(std::mem::take(bytes))),
            WireBody::Stream(stream) => poll_fn(|cx| stream.as_mut().poll_next(cx)).await,
        }
    }

    pub async fn collect_to(self, ceiling: usize) -> Result<Bytes, CoreError> {
        match self {
            WireBody::Bytes(bytes) => {
                if bytes.len() > ceiling {
                    return Err(CoreError::BodyTooLarge(ceiling));
                }
                Ok(bytes)
            }
            WireBody::Stream(mut stream) => {
                let mut collected = BytesMut::new();
                loop {
                    match poll_fn(|cx| stream.as_mut().poll_next(cx)).await {
                        Some(Ok(chunk)) => {
                            if collected.len() + chunk.len() > ceiling {
                                return Err(CoreError::BodyTooLarge(ceiling));
                            }
                            collected.extend_from_slice(&chunk);
                        }
                        Some(Err(error)) => return Err(error),
                        None => return Ok(collected.freeze()),
                    }
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn chunked(chunks: Vec<Result<Bytes, CoreError>>) -> WireBody {
        WireBody::Stream(Box::pin(futures_util::stream::iter(chunks)))
    }

    #[test]
    fn collected_bytes_below_the_ceiling_are_identity() {
        let body = WireBody::Bytes(Bytes::from_static(b"hello"));
        let collected = futures_executor::block_on(body.collect_to(16)).unwrap();
        assert_eq!(collected, Bytes::from_static(b"hello"));
    }

    #[test]
    fn a_stream_collects_chunks_in_order_below_the_ceiling() {
        let body = chunked(vec![
            Ok(Bytes::from_static(b"hel")),
            Ok(Bytes::from_static(b"lo")),
        ]);
        let collected = futures_executor::block_on(body.collect_to(16)).unwrap();
        assert_eq!(collected, Bytes::from_static(b"hello"));
    }

    #[test]
    fn unary_and_streaming_bodies_share_the_chunk_api() {
        let mut unary = WireBody::Bytes(Bytes::from_static(b"one"));
        assert_eq!(
            futures_executor::block_on(unary.next_chunk())
                .unwrap()
                .unwrap(),
            Bytes::from_static(b"one")
        );
        assert!(futures_executor::block_on(unary.next_chunk()).is_none());

        let mut streaming = chunked(vec![
            Ok(Bytes::from_static(b"two")),
            Ok(Bytes::from_static(b"three")),
        ]);
        assert_eq!(
            futures_executor::block_on(streaming.next_chunk())
                .unwrap()
                .unwrap(),
            Bytes::from_static(b"two")
        );
        assert_eq!(
            futures_executor::block_on(streaming.next_chunk())
                .unwrap()
                .unwrap(),
            Bytes::from_static(b"three")
        );
        assert!(futures_executor::block_on(streaming.next_chunk()).is_none());
    }

    #[test]
    fn collection_above_the_ceiling_fails_body_too_large() {
        let oversized = WireBody::Bytes(Bytes::from_static(b"toolarge"));
        assert!(matches!(
            futures_executor::block_on(oversized.collect_to(4)),
            Err(CoreError::BodyTooLarge(4))
        ));
        let streamed = chunked(vec![
            Ok(Bytes::from_static(b"too")),
            Ok(Bytes::from_static(b"large")),
        ]);
        assert!(matches!(
            futures_executor::block_on(streamed.collect_to(4)),
            Err(CoreError::BodyTooLarge(4))
        ));
    }
}