use super::{stream::FormatErrorInner, unfiltering_buffer::UnfilteringBuffer, DecodingError};

use fdeflate::Decompressor;

/// An inplace buffer for decompression and filtering of PNG rowlines.

///

/// The underlying data structure is a vector, with additional markers denoting a region of bytes

/// that are utilized by the decompression but not yet available to arbitrary modifications. The

/// caller can still shift around data between calls to the stream decompressor as long as the data

/// in the marked region is not modified and the indices adjusted accordingly. See

/// [`UnfilterRegion`] that contains these markers.

///

/// Violating the invariants, i.e. modifying bytes in the marked region, results in absurdly wacky

/// decompression output or panics but not undefined behavior.

pub struct UnfilterBuf<'data> {

    /// The data container. Starts with arbitrary data unrelated to the decoder, a slice of decoder

    /// private data followed by free space for further decoder output. The regions are delimited

    /// by `filled` and `available` which must be updated accordingly.

    pub(crate) buffer: &'data mut Vec<u8>,

    /// Where we record changes to the out position.

    pub(crate) filled: &'data mut usize,

    /// Where we record changes to the available byte.

    pub(crate) available: &'data mut usize,

}

/// A region into a buffer utilized as a [`UnfilterBuf`].

///

/// The span of data denoted by `filled..available` is the region of bytes that must be preserved

/// for use by the decompression algorithm. It may be moved, e.g. by subtracting the same amount

/// from both of these fields. Always ensure that `filled <= available`, the library does not

/// violate this invariant when modifying this struct as an [`UnfilterBuf`].

#[derive(Default, Clone, Copy)]

pub struct UnfilterRegion {

    /// The past-the-end index of byte that are allowed to be modified.

    pub available: usize,

    /// The past-the-end of bytes that have been written to.

    pub filled: usize,

}

/// Ergonomics wrapper around `miniz_oxide::inflate::stream` for zlib compressed data.

pub(super) struct ZlibStream {

    /// Current decoding state.

    state: Box<fdeflate::Decompressor>,

    /// If there has been a call to decompress already.

    started: bool,

    /// Ignore and do not calculate the Adler-32 checksum. Defaults to `true`.

    ///

    /// This flag overrides `TINFL_FLAG_COMPUTE_ADLER32`.

    ///

    /// This flag should not be modified after decompression has started.

    ignore_adler32: bool,

}

impl ZlibStream {

    // [PNG spec](https://www.w3.org/TR/2003/REC-PNG-20031110/#10Compression) says that

    // "deflate/inflate compression with a sliding window (which is an upper bound on the

    // distances appearing in the deflate stream) of at most 32768 bytes".

    //

    // `fdeflate` requires that we keep this many most recently decompressed bytes in the

    // `out_buffer` - this allows referring back to them when handling "length and distance

    // codes" in the deflate stream).

    const LOOKBACK_SIZE: usize = 32768;

    pub(crate) fn new() -> Self {

        ZlibStream {

            state: Box::new(Decompressor::new()),

            started: false,

            ignore_adler32: true,

        }

    }

    pub(crate) fn reset(&mut self) {

        self.started = false;

        *self.state = Decompressor::new();

    }

    /// Set the `ignore_adler32` flag and return `true` if the flag was

    /// successfully set.

    ///

    /// The default is `true`.

    ///

    /// This flag cannot be modified after decompression has started until the

    /// [ZlibStream] is reset.

    pub(crate) fn set_ignore_adler32(&mut self, flag: bool) -> bool {

        if !self.started {

            self.ignore_adler32 = flag;

            true

        } else {

            false

        }

    }

    /// Return the `ignore_adler32` flag.

    pub(crate) fn ignore_adler32(&self) -> bool {

        self.ignore_adler32

    }

    /// Fill the decoded buffer as far as possible from `data`.

    /// On success returns the number of consumed input bytes.

    pub(crate) fn decompress(

        &mut self,

        data: &[u8],

        image_data: &mut UnfilterBuf<'_>,

    ) -> Result<usize, DecodingError> {

        // There may be more data past the adler32 checksum at the end of the deflate stream. We

        // match libpng's default behavior and ignore any trailing data. In the future we may want

        // to add a flag to control this behavior.

        if self.state.is_done() {

            return Ok(data.len());

        }

        if !self.started && self.ignore_adler32 {

            self.state.ignore_adler32();

        }

        let (buffer, filled) = image_data.borrow_mut();

        let output_limit = (filled + UnfilteringBuffer::GROWTH_BYTES).min(buffer.len());

        let (in_consumed, out_consumed) = self

            .state

            .read(data, &mut buffer[..output_limit], filled, false)

            .map_err(|err| {

                DecodingError::Format(FormatErrorInner::CorruptFlateStream { err }.into())

            })?;

        self.started = true;

        let filled = filled + out_consumed;

        image_data.filled(filled);

        if self.state.is_done() {

            image_data.commit(filled);

        } else {

            // See [`Self::LOOKBACK_SIZE`].

            image_data.commit(filled.saturating_sub(Self::LOOKBACK_SIZE));

        }

        Ok(in_consumed)

    }

    /// Called after all consecutive IDAT chunks were handled.

    ///

    /// The compressed stream can be split on arbitrary byte boundaries. This enables some cleanup

    /// within the decompressor and flushing additional data which may have been kept back in case

    /// more data were passed to it.

    pub(crate) fn finish_compressed_chunks(

        &mut self,

        image_data: &mut UnfilterBuf<'_>,

    ) -> Result<(), DecodingError> {

        if !self.started {

            return Ok(());

        }

        if self.state.is_done() {

            // We can end up here only after the [`decompress`] call above has detected the state

            // to be done, too. In this case the filled and committed amount of data are already

            // equal to each other. So neither of them needs to be touched in any way.

            return Ok(());

        }

        let (_, mut filled) = image_data.borrow_mut();

        while !self.state.is_done() {

            let (buffer, _) = image_data.borrow_mut();

            let (_in_consumed, out_consumed) =

                self.state.read(&[], buffer, filled, true).map_err(|err| {

                    DecodingError::Format(FormatErrorInner::CorruptFlateStream { err }.into())

                })?;

            filled += out_consumed;

            if !self.state.is_done() {

                image_data.flush_allocate();

            }

        }

        image_data.filled(filled);

        image_data.commit(filled);

        Ok(())

    }

}

impl UnfilterRegion {

    /// Use this region to decompress new filtered rowline data.

    ///

    /// Pass the wrapped buffer to

    /// [`StreamingDecoder::update`][`super::stream::StreamingDecoder::update`] to fill it with

    /// data and update the region indices.

    pub fn as_buf<'data>(&'data mut self, buffer: &'data mut Vec<u8>) -> UnfilterBuf<'data> {

        UnfilterBuf {

            buffer,

            filled: &mut self.filled,

            available: &mut self.available,

        }

    }

}

impl UnfilterBuf<'_> {

    pub(crate) fn borrow_mut(&mut self) -> (&mut [u8], usize) {

        (self.buffer, *self.filled)

    }

    pub(crate) fn filled(&mut self, filled: usize) {

        *self.filled = filled;

    }

    pub(crate) fn commit(&mut self, howmany: usize) {

        *self.available = howmany;

    }

    pub(crate) fn flush_allocate(&mut self) {

        let len = self.buffer.len() + 32 * 1024;

        self.buffer.resize(len, 0);

    }

}