serialqoi 0.2.0

use std::io::Write;

use anyhow::{anyhow, Context, Error};

use crate::common::{hash, Channel, Colorspace, Rgb, Rgba};

/// Encodes a qoi image
///
/// This encodes an image writing it to a [`Write`r](std::io::Write).
///
/// If you're writing to a [`File`](std::fs::File) or similar you
/// probably want to wrap it in a [`BufWriter`](std::io::BufWriter)
/// for performance reasons.
#[derive(Clone, Debug)]
pub struct QoiWriter<W: Write>
{
    writer: W,
    width: u32,
    height: u32,
    channels: Channel,
    colorspace: Colorspace,
    previous_color: Rgba,
    index: [Rgba; 64],
    pixels_seen: u64,
    current_run_length: u8,
    finished: bool,
}

impl<W: Write> Drop for QoiWriter<W>
{
    fn drop(&mut self)
    {
        // Errors sadly cannot be handled here.
        drop(self.close_inner());
    }
}

impl<W: Write> QoiWriter<W>
{
    /// Creates a new [`QoiWriter`]
    ///
    /// This function creates a new [`QoiWriter`] from a
    /// [`Write`r](std::io::Write).
    ///
    /// The width, height, number of channels and the colorspace has
    /// to be given so that the correct header can be written.  The
    /// last two of these don't change the encoding in any way (which
    /// also means that theoretically even if formally there is no
    /// alpha channel you could write alpha values.  It should be
    /// obvious that you *really* shouldn't do that since it's not
    /// standard conform).
    ///
    /// # Errors
    /// An error is returned if the header couldn't be written.
    pub fn new(
        mut writer: W,
        width: u32,
        height: u32,
        channels: Channel,
        colorspace: Colorspace,
    ) -> Result<Self, Error>
    {
        let mut header = Vec::with_capacity(14);
        header.extend(b"qoif");
        header.extend(width.to_be_bytes());
        header.extend(height.to_be_bytes());

        match channels
        {
            Channel::Rgb => header.push(3),
            Channel::Rgba => header.push(4),
        }

        match colorspace
        {
            Colorspace::Srgb => header.push(0),
            Colorspace::Linear => header.push(1),
        }

        writer.write_all(&header).context("Couldn't write header")?;

        Ok(Self {
            writer,
            width,
            height,
            channels,
            colorspace,
            previous_color: Rgba::new(0, 0, 0, 255),
            index: [Rgba::new(0, 0, 0, 0); 64],
            pixels_seen: 0,
            current_run_length: 0,
            finished: false,
        })
    }

    /// Writes a single pixel (with alpha)
    ///
    /// Writes a single pixel or buffers it if this is necessary for
    /// compression.  Due to this buffering you have to
    /// [`flush()`](Self::flush), [`close()`](Self::close) or
    /// [`drop()`](std::mem::drop) the image to have it written
    /// immediately.
    ///
    /// # Errors
    /// A error is returned if already all pixels were written or if
    /// the underlying writer returned one.
    pub fn write_rgba(&mut self, color: Rgba) -> Result<(), Error>
    {
        self.write_inner(color).context("Error in writing pixel")?;

        if self.width as u64 * self.height as u64 == self.pixels_seen
        {
            self.finish().context("Error in finishing image")?;
        }

        self.previous_color = color;
        self.pixels_seen += 1;
        self.index[hash(color)] = color;

        Ok(())
    }

    /// Writes a single pixel (without alpha)
    ///
    /// Writes a single pixel or buffers it if this is necessary for
    /// compression.  Due to this buffering you have to
    /// [`flush()`](Self::flush), [`close()`](Self::close) or
    /// [`drop()`](std::mem::drop) the image to have it written
    /// immediately.
    ///
    /// The alpha value is taken from the previous pixel.
    ///
    /// # Errors
    /// A error is returned if already all pixels were written or if
    /// the underlying writer returned one.
    pub fn write_rgb(&mut self, color: Rgb) -> Result<(), Error>
    {
        let Rgb { r, g, b } = color;
        let a = self.previous_color.a;

        // No `.context(…)` here, since this couldn't add any
        // additional information.
        self.write_rgba(Rgba { r, g, b, a })
    }

    /// Writes multiple pixels (with alpha)
    ///
    /// Writes all the pixels in the [Iterator] or buffers them if
    /// this is necessary for compression.  Due to this buffering you
    /// have to [`flush()`](Self::flush), [`close()`](Self::close) or
    /// [`drop()`](std::mem::drop) the image to have it written
    /// immediately.
    ///
    /// # Errors
    /// A error is returned if already all pixels were written or if
    /// the underlying writer returned one.
    pub fn write_rgbas<I>(&mut self, pixels: I) -> Result<(), Error>
    where
        I: IntoIterator<Item = Rgba>,
    {
        pixels
            .into_iter()
            .try_for_each(|pixel| self.write_rgba(pixel))
    }

    /// Writes multiple pixels (without alpha)
    ///
    /// Writes all the pixels in the [Iterator] or buffers them if
    /// this is necessary for compression.  Due to this buffering you
    /// have to [`flush()`](Self::flush), [`close()`](Self::close) or
    /// [`drop()`](std::mem::drop) the image to have it written
    /// immediately.
    ///
    /// The alpha value is taken from the previous pixel.
    ///
    /// # Errors
    /// A error is returned if already all pixels were written or if
    /// the underlying writer returned one.
    pub fn write_rgbs<I>(&mut self, pixels: I) -> Result<(), Error>
    where
        I: IntoIterator<Item = Rgb>,
    {
        pixels
            .into_iter()
            .try_for_each(|pixel| self.write_rgb(pixel))
    }

    /// Flushes the image
    ///
    /// Flushes the image and the underlying writer.  This should
    /// never be necessary since it's done automatically when
    /// [`closed`](Self::close) or [`dropped`](std::mem::drop).
    ///
    /// # Errors
    /// A error is returned if the underlying writer returned an
    /// error.
    pub fn flush(&mut self) -> Result<(), Error>
    {
        self.write_run_length()
            .context("Error in writing buffered pixels")?;
        self.writer
            .flush()
            .context("Error in writing buffered bytes")
    }

    /// Closes the image
    ///
    /// Closes the image and the underlying writer.  This
    /// automatically done on [`drop`](std::mem::drop) when it goes
    /// out of scope.
    ///
    /// # Errors
    /// Returns an error if the underlying writer returned one.
    pub fn close(mut self) -> Result<(), Error>
    {
        self.close_inner()
    }

    /// Returns the dimensions of the image
    ///
    /// Returns the dimensions (`(width, height)`) of the image.
    #[must_use]
    pub const fn dimensions(&self) -> (u32, u32)
    {
        (self.width, self.height)
    }

    /// Returns the width of the image
    #[must_use]
    pub const fn width(&self) -> u32
    {
        self.width
    }

    /// Returns the height of the image
    #[must_use]
    pub const fn height(&self) -> u32
    {
        self.height
    }

    /// Returns the number of channels of the image
    ///
    /// Returns whether the alpha channel is used ([`Channel::Rgba`])
    /// or not ([`Channel::Rgb`]).
    ///
    /// Even if the alpha channel is disabled, a [`Rgba`] is decoded
    /// instead of an [`Rgb`](crate::common::Rgb) (the alpha channel
    /// is then always `255`).
    ///
    /// The result of this function changes nothing in the encoding.
    #[must_use]
    pub const fn channels(&self) -> Channel
    {
        self.channels
    }

    // The "sRGB" in the second paragraph shouldn't be in backticks.
    #[allow(clippy::doc_markdown)]
    /// Returns the colorspace of the image
    ///
    /// Return whether all channels are linear
    /// ([`Colorspace::Linear`]) or if it's sRGB and only the alpha
    /// channel (if it exists) is linear ([`Colorspace::Srgb`]).
    ///
    /// The result of this function changes nothing in the encoding.
    #[must_use]
    pub const fn colorspace(&self) -> Colorspace
    {
        self.colorspace
    }

    fn finish(&mut self) -> Result<(), Error>
    {
        if self.finished
        {
            Ok(())
        }
        else
        {
            self.finished = true;
            self.writer
                .write_all(&[0, 0, 0, 0, 0, 0, 0, 1])
                .context("Error writing the finishing bytes")
        }
    }

    // This is for allowing "dr_dg" and "db_dg".  This is done due to
    // them being the official names from the standard.
    #[allow(clippy::similar_names)]
    fn write_inner(&mut self, color: Rgba) -> Result<(), Error>
    {
        if self.width as u64 * self.height as u64 <= self.pixels_seen
        {
            // TODO: Change error to ErrorKind::StorageFull once #86442 gets stabilized.
            return Err(anyhow!("all pixels are already written"));
        }

        // QOI_OP_RUN for previous pixel(s)
        if color != self.previous_color && self.current_run_length > 0
        {
            self.write_run_length()
                .context("Error writing the run length")?;
            self.current_run_length = 0;
        }

        if color.a == self.previous_color.a
        {
            // QOI_OP_RUN
            if color == self.previous_color
            {
                self.current_run_length += 1;

                if self.current_run_length == 62
                {
                    self.write_run_length()
                        .context("Error writing the full run length")?;
                    self.current_run_length = 0;
                }

                return Ok(());
            }

            // QOI_OP_INDEX
            if let Some(index) = self.index.iter().position(|&x| x == color)
            {
                return self
                    .writer
                    .write_all(&[index as u8])
                    .context("Error writing INDEX chunk");
            }

            // QOI_OP_DIFF
            let dg = color.g.wrapping_sub(self.previous_color.g).wrapping_add(2);
            let dr = color.r.wrapping_sub(self.previous_color.r).wrapping_add(2);
            let db = color.b.wrapping_sub(self.previous_color.b).wrapping_add(2);

            if dr < 4 && dg < 4 && db < 4
            {
                return self
                    .writer
                    .write_all(&[64 + dr * 16 + dg * 4 + db])
                    .context("Error writing DIFF chunk");
            }

            // QOI_OP_LUMA
            let dr_dg = dr.wrapping_sub(dg).wrapping_add(8);
            let db_dg = db.wrapping_sub(dg).wrapping_add(8);
            let dg = dg.wrapping_add(30);

            if dg < 64 && dr_dg < 16 && db_dg < 16
            {
                return self
                    .writer
                    .write_all(&[128 + dg, dr_dg * 16 + db_dg])
                    .context("Error writing LUMA chunk");
            }

            // QOI_OP_RGB
            self.writer
                .write_all(&[0xfe, color.r, color.g, color.b])
                .context("Error writing RGB chunk")
        }
        else
        {
            // QOI_OP_RGBA
            self.writer
                .write_all(&[0xff, color.r, color.g, color.b, color.a])
                .context("Error writing RGBA chunk")
        }
    }

    fn write_run_length(&mut self) -> Result<(), Error>
    {
        if self.current_run_length == 0
        {
            return Ok(());
        }

        let val = 3 * 64 + self.current_run_length - 1;
        self.current_run_length = 0;

        self.writer
            .write_all(&[val])
            .context("Error writing run length")
    }

    fn close_inner(&mut self) -> Result<(), Error>
    {
        self.write_run_length()
            .context("Error in writing buffered pixels")?;
        self.finish()
            .context("Error in finishing writing the image")?;
        self.writer.flush().context("Error in flushing the writer")
    }
}