bitbottle 0.10.0

use std::cell::RefCell;
use std::fmt;
use std::io::Read;
use std::rc::Rc;
use crate::bottle_error::BottleResult;


// Extra utility types for dealing with `io::Read` and bottles:
//
// - `ReadStream`: add a `finish` method to `io:Read` so a stream can read
//   any trailing framing bytes before handing off to an outer layer
// - `ReadStreamRef`: shared reference to `ReadStream`
// - `StreamBottle`: trait for bottles that contain an inner data stream
//   and can provide a `ReadStreamRef` to it
// - `WrappedRead`: allow a normal `io::Read` to be used wherever
//   `ReadStream` can, by giving it a no-op `finish()` method


/// Enhancement of `io::Read` that also contains a `finish()` method to drain
/// the remaining (inner) stream before moving on.
pub trait ReadStream: fmt::Debug + Read {
    fn finish(&mut self) -> BottleResult<()>;
}


/// A ref-counted cell for a `ReadStream`, so the code that reads the inner
/// stream may be separated from the original bottle reader and closed later.
#[derive(Clone, Debug)]
pub struct ReadStreamRef {
    pub streaming_bottle: Rc<RefCell<dyn ReadStream>>,
}

impl ReadStreamRef {
    pub fn new(streaming_bottle: Rc<RefCell<dyn ReadStream>>) -> ReadStreamRef {
        ReadStreamRef { streaming_bottle }
    }

    pub fn finish(&mut self) -> BottleResult<()> {
        self.streaming_bottle.borrow_mut().finish()
    }
}

impl Read for ReadStreamRef {
    fn read(&mut self, buffer: &mut [u8]) -> std::io::Result<usize> {
        self.streaming_bottle.borrow_mut().read(buffer)
    }
}

impl ReadStream for ReadStreamRef {
    fn finish(&mut self) -> BottleResult<()> {
        self.streaming_bottle.borrow_mut().finish()
    }
}


/// A bottle that can produce a `ReadStream` (`io::Read` + `finish()`) of its
/// inner contents. This new stream consumes the original bottle reader, and
/// the stream will read to the end of its inner data (consuming trailing
/// framing bytes) only when you call `finish()`.
pub trait StreamBottle {
    /// turn myself into a cell, wrap that in an object that implements Read,
    /// and wrap THAT in a cell that implements Read. the end result is an
    /// opaque heap-managed object that appears to be a readable byte stream,
    /// pulling data out of this bottle.
    fn stream(self) -> ReadStreamRef;
}


/// A wrapped `io::Read` that behaves like a `ReadStream` so it can be used
/// in the same places as an inner stream.
pub struct WrappedRead<R: Read> {
    reader: R
}

impl<R: Read> WrappedRead<R> {
    pub fn new(reader: R) -> WrappedRead<R> {
        WrappedRead { reader }
    }
}

impl<R: Read> fmt::Debug for WrappedRead<R> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "WrappedRead(<io::Read>)")
    }
}

impl<R: Read + 'static> StreamBottle for WrappedRead<R> {
    fn stream(self) -> ReadStreamRef {
        ReadStreamRef::new(Rc::new(RefCell::new(self)))
    }
}

impl<R: Read> ReadStream for WrappedRead<R> {
    fn finish(&mut self) -> BottleResult<()> {
        Ok(())
    }
}

impl<R: Read> Read for WrappedRead<R> {
    fn read(&mut self, buffer: &mut [u8]) -> std::io::Result<usize> {
        self.reader.read(buffer)
    }
}


/// Read and discard data from an `io::Read`. Returns the size of the stream,
/// and if the size was less than `max_capture` bytes, also returns the data
/// from the stream.
pub fn drain_stream<R: Read>(mut reader: R, max_capture: usize) -> std::io::Result<(usize, Option<Vec<u8>>)> {
    let mut byte_count = 0;
    let mut buffer = vec![0u8; 0x1_0000];
    let mut captured = Vec::new();

    loop {
        let n = reader.read(&mut buffer)?;
        if n == 0 {
            return Ok((byte_count, (byte_count <= max_capture).then_some(captured)));
        }
        byte_count += n;
        if captured.len() < max_capture {
            let remain = max_capture - captured.len();
            captured.extend_from_slice(&buffer[.. std::cmp::max(remain, n)]);
        }
    }
}