bitbottle 0.9.1

use std::cell::RefCell;
use std::io::{Read, Write};
use std::rc::Rc;
use crate::bottle::{BottleReader, BottleStream, BottleWriter};
use crate::bottle_error::{BottleError, BottleResult};
use crate::bottle_cap::BottleType;
use crate::compression::{CompressionAlgorithm, Compressor, Decompressor};
use crate::header::Header;

const FIELD_COMPRESSION_TYPE_INT: u8 = 0;

const DEFAULT_BLOCK_SIZE_BITS: usize = 20;


/// A `Write` that writes all data into a bottle of type `Compressed`.
pub struct CompressedBottleWriter<W: Write> {
    compressor: Compressor<BottleWriter<W>>,
}

impl<W: Write> CompressedBottleWriter<W> {
    pub fn new(writer: W, algorithm: CompressionAlgorithm) -> BottleResult<CompressedBottleWriter<W>> {
        let mut header = Header::new();
        header.add_int(FIELD_COMPRESSION_TYPE_INT, (algorithm as u8) as u64)?;
        let mut bottle_writer = BottleWriter::new(writer, BottleType::Compressed, header, DEFAULT_BLOCK_SIZE_BITS)?;

        bottle_writer.write_data_stream()?;
        let compressor = Compressor::new(bottle_writer, algorithm)?;
        Ok(CompressedBottleWriter { compressor })
    }

    pub fn close(self) -> BottleResult<W> {
        let mut bottle_writer = self.compressor.close()?;
        bottle_writer.close_stream()?;
        bottle_writer.close()
    }
}

impl<W: Write> Write for CompressedBottleWriter<W> {
    fn write(&mut self, data: &[u8]) -> std::io::Result<usize> {
        self.compressor.write(data)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        self.compressor.flush()
    }
}


// snap's Decoder doesn't implement `into_inner()` so we have to pull some
// wacky antics to get our original `Read` back when we're done:
//
// - `BottleReader` in an `Rc<RefCell<_>>` so we can borrow it temporarily on
//   each `read()`, then take it for real when the bottle is closed
// - `CompressedInner` to hold the other `BottleReader` reference and wrap
//   `read()` for snap, by using the reference just long enough to handle the
//   read
// - store the snap Decoder in an option, so we can drop it early during
//   `close()`, reducing the `BottleReader` references back to one, letting
//   us dissolve it back into the original `Read`


/// Given a `BottleReader` for a bottle of type `Compressed`, decompress the
/// inner data bottle and provide it in the form of a `Read`.
pub struct CompressedBottleReader<R: Read> {
    pub algorithm: CompressionAlgorithm,
    bottle_reader: Rc<RefCell<BottleReader<R>>>,
    decompressor: Option<Decompressor<CompressedInner<R>>>,
}

// sneaky way to get a Read implementation for the bottle reader's data stream
struct CompressedInner<R: Read> {
    bottle_reader: Rc<RefCell<BottleReader<R>>>,
}

impl<R: Read> CompressedBottleReader<R> {
    pub fn new(mut bottle_reader: BottleReader<R>) -> BottleResult<CompressedBottleReader<R>> {
        let cap = &bottle_reader.bottle_cap;
        if cap.bottle_type != BottleType::Compressed {
            return Err(BottleError::WrongBottleType { expected: BottleType::Compressed, got: cap.bottle_type });
        }

        let id = cap.header.get_int(FIELD_COMPRESSION_TYPE_INT).ok_or(BottleError::UnknownCompression)?;
        let algorithm: CompressionAlgorithm = (id as u8).try_into().map_err(|_| BottleError::UnknownCompression)?;

        if bottle_reader.next_stream()? != BottleStream::Data {
            return Err(BottleError::WrongStreamType);
        }

        let bottle_reader = Rc::new(RefCell::new(bottle_reader));
        let inner = CompressedInner { bottle_reader: Rc::clone(&bottle_reader) };
        let decompressor = Some(Decompressor::new(inner, algorithm)?);
        Ok(CompressedBottleReader { algorithm, bottle_reader, decompressor })
    }

    pub fn close(mut self) -> BottleResult<R> {
        // force the lingering reference to bottle_reader to drop:
        self.decompressor.take().ok_or(BottleError::InvalidBottleState)?;

        let mut bottle_reader = Rc::try_unwrap(self.bottle_reader).map_err(|_| {
            BottleError::InvalidBottleState
        })?.into_inner();
        bottle_reader.close_stream()?;
        if bottle_reader.next_stream()? != BottleStream::End {
            return Err(BottleError::InvalidBottleState);
        }
        bottle_reader.close()
    }
}

impl<R: Read> Read for CompressedBottleReader<R> {
    fn read(&mut self, buffer: &mut [u8]) -> std::io::Result<usize> {
        self.decompressor.as_mut().ok_or_else(|| BottleError::InvalidBottleState.to_io_error())?.read(buffer)
    }
}

impl<R: Read> Read for CompressedInner<R> {
    fn read(&mut self, buffer: &mut [u8]) -> std::io::Result<usize> {
        self.bottle_reader.borrow_mut().data_stream().map_err(|e| e.to_io_error())?.read(buffer)
    }
}


#[cfg(test)]
mod test {
    use std::io::{Read, Write};
    use crate::bottle::BottleReader;
    use crate::bottle_cap::BottleType;
    use crate::compressed_bottle::CompressedBottleReader;
    use crate::compression::CompressionAlgorithm;
    use super::CompressedBottleWriter;

    #[test]
    fn snappy() {
        let mut buffer = Vec::new();
        let data = vec![0xaa; 1024];
        {
            let mut bottle_writer = CompressedBottleWriter::new(&mut buffer, CompressionAlgorithm::SNAPPY).unwrap();
            bottle_writer.write_all(&data).unwrap();
            bottle_writer.close().unwrap();
        }

        assert!(buffer.len() < data.len());

        let bottle_reader = BottleReader::new(&buffer[..]).unwrap();
        assert_eq!(bottle_reader.bottle_cap.bottle_type, BottleType::Compressed);
        {
            let mut compressed_reader = CompressedBottleReader::new(bottle_reader).unwrap();
            let mut data_buffer = Vec::new();
            let len = compressed_reader.read_to_end(&mut data_buffer).unwrap();
            compressed_reader.close().unwrap();
            assert_eq!(&data_buffer[0 .. len], data);
        }
    }

    #[test]
    fn lzma2() {
        let mut buffer = Vec::new();
        let data = vec![0xaa; 1024];
        {
            let mut bottle_writer = CompressedBottleWriter::new(&mut buffer, CompressionAlgorithm::LZMA2).unwrap();
            bottle_writer.write_all(&data).unwrap();
            bottle_writer.close().unwrap();
        }

        assert!(buffer.len() < data.len());

        let bottle_reader = BottleReader::new(&buffer[..]).unwrap();
        assert_eq!(bottle_reader.bottle_cap.bottle_type, BottleType::Compressed);
        {
            let mut compressed_reader = CompressedBottleReader::new(bottle_reader).unwrap();
            let mut data_buffer = Vec::new();
            let len = compressed_reader.read_to_end(&mut data_buffer).unwrap();
            compressed_reader.close().unwrap();
            assert_eq!(&data_buffer[0 .. len], data);
        }
    }
}