serialqoi 0.1.7

use std::io::{Error, ErrorKind, Read, Result};

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

/// Decodes a qoi image
///
/// This decodes an image reading it from a [`Read`er](std::io::Read).
///
/// If you're reading from a [`File`](std::fs::File) or similar you
/// probably want to wrap it in a [`BufReader`](std::io::BufReader)
/// for performance reasons.
#[derive(Copy, Clone, Debug)]
pub struct QoiReader<R: Read>
{
    reader: R,
    width: u32,
    height: u32,
    channels: Channel,
    colorspace: Colorspace,
    previous_color: Rgba,
    index: [Rgba; 64],
    pixels_seen: u64,
    remaining_run_length: u8,
    fused: bool,
}

impl<R: Read> QoiReader<R>
{
    /// Creates a new [`QoiReader`]
    ///
    /// This function creates a new [`QoiReader`] from a
    /// [`Read`er](std::io::Read).
    ///
    /// # Errors
    /// An error is returned if the header couldn't be read or
    /// correctly parsed.
    pub fn new(mut reader: R) -> Result<Self>
    {
        let mut buf = [0; 4];

        reader.read_exact(&mut buf)?;
        if buf != [b'q', b'o', b'i', b'f']
        {
            return Err(Error::new(ErrorKind::InvalidData, "wrong magic bytes"));
        }

        reader.read_exact(&mut buf)?;
        let width = u32::from_be_bytes(buf);
        reader.read_exact(&mut buf)?;
        let height = u32::from_be_bytes(buf);

        reader.read_exact(&mut buf[..1])?;
        let channels = match buf[0]
        {
            3 => Channel::Rgb,
            4 => Channel::Rgba,
            _ =>
            {
                return Err(Error::new(
                    ErrorKind::InvalidData,
                    "wrong number of channels",
                ))
            }
        };

        reader.read_exact(&mut buf[..1])?;
        let colorspace = match buf[0]
        {
            0 => Colorspace::Srgb,
            1 => Colorspace::Linear,
            _ => return Err(Error::new(ErrorKind::InvalidData, "invalid colorspace")),
        };

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

    /// 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
    }
}

impl<R: Read> Iterator for QoiReader<R>
{
    type Item = Result<Rgba>;

    // 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 next(&mut self) -> Option<Self::Item>
    {
        macro_rules! handle {
            ($val: expr) => {{
                match $val
                {
                    Ok(x) => x,
                    Err(err) =>
                    {
                        self.fused = true;
                        return Some(Err(err));
                    }
                }
            }};
        }

        if self.fused
        {
            return None;
        }
        if self.width as u64 * self.height as u64 <= self.pixels_seen
        {
            self.fused = true;
            return None;
        }
        if self.remaining_run_length > 0
        {
            self.remaining_run_length -= 1;
            self.pixels_seen += 1;
            return Some(Ok(self.previous_color));
        }

        let mut buf = [0; 5];

        handle!(self.reader.read_exact(&mut buf[..1]));

        let tag = buf[0] / 64;
        let color;

        if buf[0] >= 0xfe
        {
            let end_index = buf[0] as usize - 0xfe + 4;
            handle!(self.reader.read_exact(&mut buf[1..end_index]));

            if buf[0] == 0xff
            {
                color = Rgba::new(buf[1], buf[2], buf[3], buf[4]);
            }
            else
            {
                color = Rgba::new(buf[1], buf[2], buf[3], self.previous_color.a);
            }
        }
        else if tag == 0b00
        {
            color = self.index[buf[(tag % 0b11_1111) as usize] as usize];
        }
        else if tag == 0b01
        {
            let dr = ((buf[0] / 4) % 4).wrapping_sub(2);
            let dg = ((buf[0] / 2) % 4).wrapping_sub(2);
            let db = (buf[0] % 4).wrapping_sub(2);

            let Rgba { r, g, b, a } = self.previous_color;

            let r = r.wrapping_add(dr);
            let g = g.wrapping_add(dg);
            let b = b.wrapping_add(db);

            color = Rgba { r, g, b, a };
        }
        else if tag == 0b10
        {
            handle!(self.reader.read_exact(&mut buf[1..2]));

            let dg = (buf[0] % 64).wrapping_sub(32);
            let dr_dg = (buf[1] / 16).wrapping_sub(8);
            let db_dg = (buf[1] % 16).wrapping_sub(8);

            let Rgba { r, g, b, a } = self.previous_color;

            let r = r.wrapping_add(dg).wrapping_add(dr_dg);
            let g = g.wrapping_add(dg);
            let b = b.wrapping_add(dg).wrapping_add(db_dg);

            color = Rgba { r, g, b, a };
        }
        else if tag == 0b11
        {
            self.remaining_run_length = buf[0] % 64;

            color = self.previous_color;
        }
        else
        {
            unreachable!("This shouldn't be reachable!");
        }

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

        Some(Ok(color))
    }
}