datum-core 0.10.1

use crate::stream::{BoxStream, Flow};
use crate::{StreamError, StreamResult};
use flate2::Compression as FlateCompression;
use flate2::bufread::GzDecoder as GzBufDecoder;
use flate2::write::{GzEncoder, ZlibEncoder};
use flate2::{Decompress, FlushDecompress, Status};
use std::collections::VecDeque;
use std::io::{self, BufRead, Read, Write};

const DECOMPRESS_CHUNK_SIZE: usize = 8192;
// A decompressor pull may prefetch a small bounded batch to amortize codec calls,
// but never queues an input chunk's full expansion.
const DECOMPRESS_OUTPUT_BUDGET: usize = DECOMPRESS_CHUNK_SIZE * 16;

#[derive(Clone)]
enum Terminal {
    Complete,
    Error(StreamError),
}

fn sticky_terminal<T>(terminal: &Terminal) -> Option<StreamResult<T>> {
    match terminal {
        Terminal::Complete => None,
        Terminal::Error(error) => Some(Err(error.clone())),
    }
}

fn codec_error<E: std::fmt::Display>(error: E) -> StreamError {
    StreamError::Failed(error.to_string())
}

/// In-process byte-stream (de)compression flows backed by `flate2`. Mirrors Akka's `Compression`.
///
/// Each flow consumes and produces `Vec<u8>` chunks; output chunk boundaries do not line up with
/// input chunk boundaries. A truncated compressed input fails the decompressing flows with
/// [`StreamError`].
pub struct Compression;

impl Compression {
    /// Compresses the byte stream in the gzip format.
    #[must_use]
    pub fn gzip() -> Flow<Vec<u8>, Vec<u8>> {
        Flow::from_transform(|input| Box::new(CompressStream::gzip(input)) as BoxStream<Vec<u8>>)
    }

    /// Compresses the byte stream in the zlib/deflate format.
    #[must_use]
    pub fn deflate() -> Flow<Vec<u8>, Vec<u8>> {
        Flow::from_transform(|input| Box::new(CompressStream::deflate(input)) as BoxStream<Vec<u8>>)
    }

    /// Decompresses a gzip-format byte stream (inverse of [`Compression::gzip`]).
    #[must_use]
    pub fn gunzip() -> Flow<Vec<u8>, Vec<u8>> {
        Flow::from_transform(|input| Box::new(GunzipStream::new(input)) as BoxStream<Vec<u8>>)
    }

    /// Decompresses a zlib/deflate-format byte stream (inverse of [`Compression::deflate`]).
    #[must_use]
    pub fn inflate() -> Flow<Vec<u8>, Vec<u8>> {
        Flow::from_transform(|input| Box::new(InflateStream::new(input)) as BoxStream<Vec<u8>>)
    }
}

enum EncoderKind {
    Gzip(GzEncoder<Vec<u8>>),
    Deflate(ZlibEncoder<Vec<u8>>),
}

impl EncoderKind {
    fn write_all(&mut self, chunk: &[u8]) -> std::io::Result<()> {
        match self {
            Self::Gzip(codec) => codec.write_all(chunk),
            Self::Deflate(codec) => codec.write_all(chunk),
        }
    }

    fn try_finish(&mut self) -> std::io::Result<()> {
        match self {
            Self::Gzip(codec) => codec.try_finish(),
            Self::Deflate(codec) => codec.try_finish(),
        }
    }

    fn take_output(&mut self) -> Vec<u8> {
        match self {
            Self::Gzip(codec) => std::mem::take(codec.get_mut()),
            Self::Deflate(codec) => std::mem::take(codec.get_mut()),
        }
    }
}

struct CompressStream {
    input: BoxStream<Vec<u8>>,
    codec: EncoderKind,
    pending: VecDeque<Vec<u8>>,
    finished: bool,
    terminal: Option<Terminal>,
}

impl CompressStream {
    fn gzip(input: BoxStream<Vec<u8>>) -> Self {
        Self {
            input,
            codec: EncoderKind::Gzip(GzEncoder::new(Vec::new(), FlateCompression::default())),
            pending: VecDeque::new(),
            finished: false,
            terminal: None,
        }
    }

    fn deflate(input: BoxStream<Vec<u8>>) -> Self {
        Self {
            input,
            codec: EncoderKind::Deflate(ZlibEncoder::new(Vec::new(), FlateCompression::default())),
            pending: VecDeque::new(),
            finished: false,
            terminal: None,
        }
    }

    fn fail<T>(&mut self, error: StreamError) -> Option<StreamResult<T>> {
        self.terminal = Some(Terminal::Error(error.clone()));
        Some(Err(error))
    }

    fn harvest_output(&mut self) {
        let output = self.codec.take_output();
        if !output.is_empty() {
            self.pending.push_back(output);
        }
    }
}

impl Iterator for CompressStream {
    type Item = StreamResult<Vec<u8>>;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(chunk) = self.pending.pop_front() {
            return Some(Ok(chunk));
        }
        if let Some(terminal) = &self.terminal {
            return sticky_terminal(terminal);
        }

        loop {
            if self.finished {
                self.terminal = Some(Terminal::Complete);
                return None;
            }

            match self.input.next() {
                Some(Ok(chunk)) => {
                    if let Err(error) = self.codec.write_all(&chunk).map_err(codec_error) {
                        return self.fail(error);
                    }
                    self.harvest_output();
                    if let Some(chunk) = self.pending.pop_front() {
                        return Some(Ok(chunk));
                    }
                }
                Some(Err(error)) => {
                    self.terminal = Some(Terminal::Error(error.clone()));
                    return Some(Err(error));
                }
                None => {
                    if let Err(error) = self.codec.try_finish().map_err(codec_error) {
                        return self.fail(error);
                    }
                    self.finished = true;
                    self.harvest_output();
                    if let Some(chunk) = self.pending.pop_front() {
                        return Some(Ok(chunk));
                    }
                }
            }
        }
    }
}

struct ChunkReader {
    input: BoxStream<Vec<u8>>,
    current: Vec<u8>,
    offset: usize,
    active: bool,
    done: bool,
    stream_error: Option<StreamError>,
}

impl ChunkReader {
    fn new(input: BoxStream<Vec<u8>>) -> Self {
        Self {
            input,
            current: Vec::new(),
            offset: 0,
            active: false,
            done: false,
            stream_error: None,
        }
    }

    fn activate(&mut self) {
        self.active = true;
    }

    fn stream_error(&self) -> Option<StreamError> {
        self.stream_error.clone()
    }
}

impl Read for ChunkReader {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let available = self.fill_buf()?;
        let count = available.len().min(buf.len());
        buf[..count].copy_from_slice(&available[..count]);
        self.consume(count);
        Ok(count)
    }
}

impl BufRead for ChunkReader {
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        if !self.active {
            return Err(io::ErrorKind::WouldBlock.into());
        }

        while self.offset >= self.current.len() && !self.done {
            self.current.clear();
            self.offset = 0;
            match self.input.next() {
                Some(Ok(chunk)) if chunk.is_empty() => {}
                Some(Ok(chunk)) => self.current = chunk,
                Some(Err(error)) => {
                    self.stream_error = Some(error.clone());
                    self.done = true;
                    return Err(io::Error::other(error.to_string()));
                }
                None => self.done = true,
            }
        }

        Ok(&self.current[self.offset..])
    }

    fn consume(&mut self, amt: usize) {
        self.offset = (self.offset + amt).min(self.current.len());
        if self.offset == self.current.len() {
            self.current.clear();
            self.offset = 0;
        }
    }
}

struct GunzipStream {
    codec: GzBufDecoder<ChunkReader>,
    finished: bool,
    terminal: Option<Terminal>,
}

impl GunzipStream {
    fn new(input: BoxStream<Vec<u8>>) -> Self {
        Self {
            codec: GzBufDecoder::new(ChunkReader::new(input)),
            finished: false,
            terminal: None,
        }
    }

    fn fail<T>(&mut self, error: StreamError) -> Option<StreamResult<T>> {
        self.terminal = Some(Terminal::Error(error.clone()));
        Some(Err(error))
    }
}

impl Iterator for GunzipStream {
    type Item = StreamResult<Vec<u8>>;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(terminal) = &self.terminal {
            return sticky_terminal(terminal);
        }
        if self.finished {
            self.terminal = Some(Terminal::Complete);
            return None;
        }

        self.codec.get_mut().activate();
        let mut output = vec![0_u8; DECOMPRESS_CHUNK_SIZE];
        match self.codec.read(&mut output) {
            Ok(0) => {
                self.finished = true;
                self.terminal = Some(Terminal::Complete);
                None
            }
            Ok(count) => {
                output.truncate(count);
                output.shrink_to_fit();
                Some(Ok(output))
            }
            Err(error) => {
                let error = self
                    .codec
                    .get_ref()
                    .stream_error()
                    .unwrap_or_else(|| codec_error(error));
                self.fail(error)
            }
        }
    }
}

struct InflateStream {
    input: BoxStream<Vec<u8>>,
    codec: Decompress,
    pending_input: Vec<u8>,
    pending_output: VecDeque<Vec<u8>>,
    pending_offset: usize,
    draining: bool,
    upstream_done: bool,
    finished: bool,
    terminal: Option<Terminal>,
}

impl InflateStream {
    fn new(input: BoxStream<Vec<u8>>) -> Self {
        Self {
            input,
            codec: Decompress::new(true),
            pending_input: Vec::new(),
            pending_output: VecDeque::new(),
            pending_offset: 0,
            draining: false,
            upstream_done: false,
            finished: false,
            terminal: None,
        }
    }

    fn fail<T>(&mut self, error: StreamError) -> Option<StreamResult<T>> {
        self.terminal = Some(Terminal::Error(error.clone()));
        Some(Err(error))
    }

    fn has_pending_input(&self) -> bool {
        self.pending_offset < self.pending_input.len()
    }

    fn load_input(&mut self, chunk: Vec<u8>) {
        self.pending_input = chunk;
        self.pending_offset = 0;
    }

    fn clear_consumed_input(&mut self) {
        if self.pending_offset >= self.pending_input.len() {
            self.pending_input.clear();
            self.pending_offset = 0;
        }
    }

    fn pump_once(&mut self, flush: FlushDecompress) -> StreamResult<Pump> {
        let before_in = self.codec.total_in();
        let before_out = self.codec.total_out();
        let mut output = vec![0_u8; DECOMPRESS_CHUNK_SIZE];
        let status = self
            .codec
            .decompress(
                &self.pending_input[self.pending_offset..],
                &mut output,
                flush,
            )
            .map_err(codec_error)?;
        let consumed = (self.codec.total_in() - before_in) as usize;
        let produced = (self.codec.total_out() - before_out) as usize;

        self.pending_offset += consumed;
        self.clear_consumed_input();

        if matches!(status, Status::StreamEnd) {
            self.finished = true;
        }

        output.truncate(produced);
        if !output.is_empty() {
            self.draining = produced == DECOMPRESS_CHUNK_SIZE && !self.finished;
            output.shrink_to_fit();
            return Ok(Pump::Output(output));
        }

        if matches!(status, Status::StreamEnd) {
            return Ok(Pump::Complete);
        }

        if consumed == 0 {
            self.draining = false;
            return Ok(Pump::NeedInput);
        }

        Ok(Pump::Continue)
    }
}

enum Pump {
    Output(Vec<u8>),
    NeedInput,
    Continue,
    Complete,
}

impl Iterator for InflateStream {
    type Item = StreamResult<Vec<u8>>;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(chunk) = self.pending_output.pop_front() {
            return Some(Ok(chunk));
        }
        if let Some(terminal) = &self.terminal {
            return sticky_terminal(terminal);
        }
        if self.finished {
            self.terminal = Some(Terminal::Complete);
            return None;
        }

        let mut produced_this_pull = 0_usize;
        let mut queued_output = false;
        loop {
            if produced_this_pull >= DECOMPRESS_OUTPUT_BUDGET {
                break;
            }
            if !self.has_pending_input() && !self.upstream_done && !self.draining {
                if queued_output {
                    break;
                }
                match self.input.next() {
                    Some(Ok(chunk)) => self.load_input(chunk),
                    Some(Err(error)) => {
                        self.terminal = Some(Terminal::Error(error.clone()));
                        return Some(Err(error));
                    }
                    None => self.upstream_done = true,
                }
            }

            let flush = if self.upstream_done {
                FlushDecompress::Finish
            } else {
                FlushDecompress::None
            };

            match self.pump_once(flush) {
                Ok(Pump::Output(chunk)) => {
                    produced_this_pull += chunk.len();
                    queued_output = true;
                    self.pending_output.push_back(chunk);
                }
                Ok(Pump::Complete) => {
                    self.terminal = Some(Terminal::Complete);
                    break;
                }
                Ok(Pump::NeedInput) => {
                    if self.upstream_done {
                        return self.fail(StreamError::Failed(
                            "truncated compressed stream".to_owned(),
                        ));
                    }
                    if queued_output {
                        break;
                    }
                }
                Ok(Pump::Continue) => {}
                Err(error) => return self.fail(error),
            }
        }

        if let Some(chunk) = self.pending_output.pop_front() {
            return Some(Ok(chunk));
        }
        if let Some(terminal) = &self.terminal {
            return sticky_terminal(terminal);
        }
        None
    }
}

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

    fn collect_chunks(flow: Flow<Vec<u8>, Vec<u8>>) -> Vec<Vec<u8>> {
        Source::from_iter([b"hello ".to_vec(), b"world".to_vec()])
            .via(flow)
            .run_with(crate::Sink::collect())
            .expect("flow materializes")
            .wait()
            .expect("flow completes")
    }

    #[test]
    fn gzip_and_gunzip_round_trip() {
        let compressed = collect_chunks(Compression::gzip());
        let decoded = Source::from_iter(compressed)
            .via(Compression::gunzip())
            .run_with(crate::Sink::collect())
            .expect("gunzip materializes")
            .wait()
            .expect("gunzip completes");

        assert_eq!(decoded.concat(), b"hello world");
    }

    #[test]
    fn deflate_and_inflate_round_trip() {
        let compressed = collect_chunks(Compression::deflate());
        let decoded = Source::from_iter(compressed)
            .via(Compression::inflate())
            .run_with(crate::Sink::collect())
            .expect("inflate materializes")
            .wait()
            .expect("inflate completes");

        assert_eq!(decoded.concat(), b"hello world");
    }

    #[test]
    fn gunzip_completes_when_output_drains_before_finish() {
        // Regression for a gunzip EOF hang: when a valid gzip stream is delivered
        // in small chunks, all decompressed output can drain before the final
        // EOF poll. The decoder must still complete instead of looping forever
        // on `input.next() == None`.
        let compressed = collect_chunks(Compression::gzip()).concat();
        let byte_chunks: Vec<Vec<u8>> = compressed.into_iter().map(|byte| vec![byte]).collect();

        let decoded = Source::from_iter(byte_chunks)
            .via(Compression::gunzip())
            .run_with(crate::Sink::collect())
            .expect("gunzip materializes")
            .wait()
            .expect("gunzip completes");

        assert_eq!(decoded.concat(), b"hello world");
    }

    #[test]
    fn inflate_completes_when_output_drains_before_finish() {
        let compressed = collect_chunks(Compression::deflate()).concat();
        let byte_chunks: Vec<Vec<u8>> = compressed.into_iter().map(|byte| vec![byte]).collect();

        let decoded = Source::from_iter(byte_chunks)
            .via(Compression::inflate())
            .run_with(crate::Sink::collect())
            .expect("inflate materializes")
            .wait()
            .expect("inflate completes");

        assert_eq!(decoded.concat(), b"hello world");
    }

    #[test]
    fn gunzip_fails_on_truncated_input() {
        let compressed = collect_chunks(Compression::gzip());
        let mut truncated = compressed.concat();
        truncated.truncate(truncated.len().saturating_sub(2));

        let result = Source::single(truncated)
            .via(Compression::gunzip())
            .run_with(crate::Sink::collect())
            .expect("gunzip materializes")
            .wait();

        assert!(matches!(result, Err(StreamError::Failed(_))));
    }

    #[test]
    fn inflate_fails_on_truncated_input() {
        let compressed = collect_chunks(Compression::deflate());
        let mut truncated = compressed.concat();
        truncated.truncate(truncated.len() / 2);

        let result = Source::single(truncated)
            .via(Compression::inflate())
            .run_with(crate::Sink::collect())
            .expect("inflate materializes")
            .wait();

        assert!(matches!(result, Err(StreamError::Failed(_))));
    }

    #[test]
    fn gunzip_emits_bounded_chunks_for_high_ratio_input() {
        let payload = vec![b'x'; DECOMPRESS_CHUNK_SIZE * 4 + 17];
        let compressed = Source::single(payload.clone())
            .via(Compression::gzip())
            .run_with(crate::Sink::collect())
            .expect("gzip materializes")
            .wait()
            .expect("gzip completes");

        let decoded = Source::single(compressed.concat())
            .via(Compression::gunzip())
            .run_with(crate::Sink::collect())
            .expect("gunzip materializes")
            .wait()
            .expect("gunzip completes");

        assert!(decoded.len() > 1);
        assert!(
            decoded
                .iter()
                .all(|chunk| chunk.len() <= DECOMPRESS_CHUNK_SIZE)
        );
        assert_eq!(decoded.concat(), payload);
    }

    #[test]
    fn inflate_emits_bounded_chunks_for_high_ratio_input() {
        let payload = vec![b'x'; DECOMPRESS_CHUNK_SIZE * 4 + 17];
        let compressed = Source::single(payload.clone())
            .via(Compression::deflate())
            .run_with(crate::Sink::collect())
            .expect("deflate materializes")
            .wait()
            .expect("deflate completes");

        let decoded = Source::single(compressed.concat())
            .via(Compression::inflate())
            .run_with(crate::Sink::collect())
            .expect("inflate materializes")
            .wait()
            .expect("inflate completes");

        assert!(decoded.len() > 1);
        assert!(
            decoded
                .iter()
                .all(|chunk| chunk.len() <= DECOMPRESS_CHUNK_SIZE)
        );
        assert_eq!(decoded.concat(), payload);
    }
}