include-pnm 1.0.0

//! Parsing of specific PNM formats.

use std::{array, fmt, num::NonZero};

use proc_macro::{Delimiter, Group, Ident, Literal, Punct, Spacing, Span, TokenStream, TokenTree};

use crate::{
    cold,
    metadata::{Depth, DepthChannels, Size},
    parse::{Header, HeaderDirective, Reader, TuplType},
};

/// Create an [`array::IntoIter`] with an easy macro.
macro_rules! array {
    ($one:expr $(,)?) => {{
        let mut x = [$one, 0, 0, 0].into_iter();
        x.next_back().unwrap();
        x.next_back().unwrap();
        x.next_back().unwrap();
        x
    }};
    ($one:expr, $two:expr $(,)?) => {{
        let mut x = [$one, $two, 0, 0].into_iter();
        x.next_back().unwrap();
        x.next_back().unwrap();
        x
    }};
    ($one:expr, $two:expr, $three:expr $(,)?) => {{
        let mut x = [$one, $two, $three, 0].into_iter();
        x.next_back().unwrap();
        x
    }};
    ($one:expr, $two:expr, $three:expr, $four:expr $(,)?) => {
        [$one, $two, $three, $four].into_iter()
    };
}

/// Makes an array with the given number of channels.
fn array_channels(arr: [u16; 4], channels: u8) -> array::IntoIter<u16, 4> {
    let mut new = arr.into_iter();
    for _ in 0..(4 - channels) {
        new.next_back();
    }
    new
}

/// Iterator converting image data.
struct Convert<I: Iterator<Item = u16>> {
    /// Initial depth/channels.
    from: Format,

    /// Final depth/channels.
    into: DepthChannels,

    /// Iterator over image data.
    data: I,
}
impl<I: Iterator<Item = u16>> Iterator for Convert<I> {
    type Item = array::IntoIter<u16, 4>;
    fn next(&mut self) -> Option<Self::Item> {
        let mut old = [0; 4];
        for channel in old.iter_mut().take(self.from.channels() as usize) {
            *channel = self.data.next()?;
        }
        Some(match (self.from, self.into) {
            (Format::Standard(DepthChannels::BlackAndWhite), DepthChannels::BlackAndWhite)
            | (
                Format::Standard(DepthChannels::BlackAndWhiteAlpha),
                DepthChannels::BlackAndWhiteAlpha,
            )
            | (Format::Standard(DepthChannels::Grayscale), DepthChannels::Grayscale)
            | (Format::Standard(DepthChannels::GrayscaleAlpha), DepthChannels::GrayscaleAlpha)
            | (Format::Standard(DepthChannels::Rgb), DepthChannels::Rgb)
            | (Format::Standard(DepthChannels::RgbAlpha), DepthChannels::RgbAlpha)
            | (Format::Standard(DepthChannels::Grayscale16), DepthChannels::Grayscale16)
            | (
                Format::Standard(DepthChannels::GrayscaleAlpha16),
                DepthChannels::GrayscaleAlpha16,
            )
            | (Format::Standard(DepthChannels::Rgb16), DepthChannels::Rgb16)
            | (Format::Standard(DepthChannels::RgbAlpha16), DepthChannels::RgbAlpha16) => {
                array_channels(old, self.into.channels() as u8)
            }
            (Format::Standard(DepthChannels::BlackAndWhite), DepthChannels::BlackAndWhiteAlpha) => {
                let old = old[0];
                array![old, 1]
            }
            (Format::Standard(DepthChannels::BlackAndWhite), DepthChannels::Grayscale) => {
                let old = old[0] * 0xFF;
                array![old]
            }
            (Format::Standard(DepthChannels::BlackAndWhite), DepthChannels::GrayscaleAlpha) => {
                let old = old[0] * 0xFF;
                array![old, 0xFF]
            }
            (Format::Standard(DepthChannels::BlackAndWhite), DepthChannels::Rgb) => {
                let old = old[0] * 0xFF;
                array![old, old, old]
            }
            (Format::Standard(DepthChannels::BlackAndWhite), DepthChannels::RgbAlpha) => {
                let old = old[0] * 0xFF;
                array![old, old, old, 0xFF]
            }
            (Format::Standard(DepthChannels::BlackAndWhite), DepthChannels::Grayscale16) => {
                let old = old[0] * 0xFFFF;
                array![old]
            }
            (Format::Standard(DepthChannels::BlackAndWhite), DepthChannels::GrayscaleAlpha16) => {
                let old = old[0] * 0xFFFF;
                array![old, 0xFFFF]
            }
            (Format::Standard(DepthChannels::BlackAndWhite), DepthChannels::Rgb16) => {
                let old = old[0] * 0xFFFF;
                array![old, old, old]
            }
            (Format::Standard(DepthChannels::BlackAndWhite), DepthChannels::RgbAlpha16) => {
                let old = old[0] * 0xFFFF;
                array![old, old, old, 0xFFFF]
            }

            (
                Format::Standard(DepthChannels::BlackAndWhiteAlpha),
                DepthChannels::GrayscaleAlpha,
            ) => {
                let alpha = old[1] * 0xFF;
                let old = old[0] * 0xFF;
                array![old, alpha]
            }
            (Format::Standard(DepthChannels::BlackAndWhiteAlpha), DepthChannels::RgbAlpha) => {
                let alpha = old[1] * 0xFF;
                let old = old[0] * 0xFF;
                array![old, old, old, alpha]
            }
            (
                Format::Standard(DepthChannels::BlackAndWhiteAlpha),
                DepthChannels::GrayscaleAlpha16,
            ) => {
                let alpha = old[1] * 0xFFFF;
                let old = old[0] * 0xFFFF;
                array![old, alpha]
            }
            (Format::Standard(DepthChannels::BlackAndWhiteAlpha), DepthChannels::RgbAlpha16) => {
                let alpha = old[1] * 0xFFFF;
                let old = old[0] * 0xFFFF;
                array![old, old, old, alpha]
            }

            (Format::Standard(DepthChannels::Grayscale), DepthChannels::GrayscaleAlpha) => {
                let old = old[0];
                array![old, 0xFF]
            }
            (Format::Standard(DepthChannels::Grayscale), DepthChannels::Rgb)
            | (Format::Standard(DepthChannels::Grayscale16), DepthChannels::Rgb16) => {
                let old = old[0];
                array![old, old, old]
            }
            (Format::Standard(DepthChannels::Grayscale), DepthChannels::RgbAlpha) => {
                let old = old[0];
                array![old, old, old, 0xFF]
            }
            (Format::Standard(DepthChannels::Grayscale), DepthChannels::Grayscale16) => {
                let old = old[0] * 0x101;
                array![old]
            }
            (Format::Standard(DepthChannels::Grayscale), DepthChannels::GrayscaleAlpha16) => {
                let old = old[0] * 0x101;
                array![old, 0xFFFF]
            }
            (Format::Standard(DepthChannels::Grayscale), DepthChannels::Rgb16) => {
                let old = old[0] * 0x101;
                array![old, old, old]
            }
            (Format::Standard(DepthChannels::Grayscale), DepthChannels::RgbAlpha16) => {
                let old = old[0] * 0x101;
                array![old, old, old, 0xFFFF]
            }

            (Format::Standard(DepthChannels::GrayscaleAlpha), DepthChannels::RgbAlpha)
            | (Format::Standard(DepthChannels::GrayscaleAlpha16), DepthChannels::RgbAlpha16) => {
                let alpha = old[1];
                let old = old[0];
                array![old, old, old, alpha]
            }
            (Format::Standard(DepthChannels::GrayscaleAlpha), DepthChannels::GrayscaleAlpha16) => {
                let alpha = old[1] * 0x101;
                let old = old[0] * 0x101;
                array![old, alpha]
            }
            (Format::Standard(DepthChannels::GrayscaleAlpha), DepthChannels::RgbAlpha16) => {
                let alpha = old[1] * 0x101;
                let old = old[0] * 0x101;
                array![old, old, old, alpha]
            }

            (Format::Standard(DepthChannels::Rgb), DepthChannels::RgbAlpha) => {
                let red = old[0];
                let green = old[1];
                let blue = old[2];
                array![red, green, blue, 0xFF]
            }
            (Format::Standard(DepthChannels::Rgb), DepthChannels::Rgb16) => {
                let red = old[0] * 0x101;
                let green = old[1] * 0x101;
                let blue = old[2] * 0x101;
                array![red, green, blue]
            }
            (Format::Standard(DepthChannels::Rgb), DepthChannels::RgbAlpha16) => {
                let red = old[0] * 0x101;
                let green = old[1] * 0x101;
                let blue = old[2] * 0x101;
                array![red, green, blue, 0xFFFF]
            }

            (Format::Standard(DepthChannels::RgbAlpha), DepthChannels::RgbAlpha16) => {
                let alpha = old[3] * 0x101;
                let red = old[0] * 0x101;
                let green = old[1] * 0x101;
                let blue = old[2] * 0x101;
                array![red, green, blue, alpha]
            }

            (Format::Standard(DepthChannels::Grayscale16), DepthChannels::GrayscaleAlpha16) => {
                let old = old[0];
                array![old, 0xFFFF]
            }
            (Format::Standard(DepthChannels::Grayscale16), DepthChannels::RgbAlpha16) => {
                let old = old[0];
                array![old, old, old, 0xFFFF]
            }

            (Format::Standard(DepthChannels::Rgb16), DepthChannels::RgbAlpha16) => {
                let red = old[0];
                let green = old[1];
                let blue = old[2];
                array![red, green, blue, 0xFFFF]
            }

            (Format::Custom { .. }, depth_channels) => {
                array_channels(old, depth_channels.channels() as u8)
            }

            _ => unreachable!("converting {:?} to {:?}", self.from, self.into),
        })
    }
}

/// Whether pixel channels should be flattened.
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
pub(crate) enum Nesting {
    /// Single-channel pixels are still emitted as arrays.
    Arrays,

    /// Single-channel pixels are inlined directly.
    Inline,
}

/// Image format.
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
pub(crate) enum Format {
    /// Standard format.
    Standard(DepthChannels),

    /// Custom format.
    Custom {
        /// Channel count.
        channels: NonZero<u8>,

        /// Maximum value.
        max_value: NonZero<u16>,
    },
}
impl Format {
    /// Create from data.
    pub(crate) fn new(channels: NonZero<u8>, max_value: NonZero<u16>) -> Format {
        (DepthChannels::new_raw(channels, max_value)).map_or(
            Format::Custom {
                channels,
                max_value,
            },
            Format::Standard,
        )
    }

    /// Approximate maximum depth.
    pub(crate) fn approx_depth(self) -> Depth {
        match self {
            Format::Standard(depth_channels) => depth_channels.depth(),
            Format::Custom { max_value, .. } => match max_value.get() {
                1 => Depth::One,
                2..=255 => Depth::Eight,
                _ => Depth::Sixteen,
            },
        }
    }

    /// Get channel count.
    pub(crate) fn channels(self) -> u8 {
        match self {
            Format::Standard(depth_channels) => depth_channels.channels() as u8,
            Format::Custom { channels, .. } => channels.get(),
        }
    }

    /// Get maximum value.
    pub(crate) fn max_value(self) -> u16 {
        match self {
            Format::Standard(depth_channels) => depth_channels.depth().max_value(),
            Format::Custom { max_value, .. } => max_value.get(),
        }
    }
}
impl fmt::Display for Format {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Format::Standard(depth_channels) => fmt::Display::fmt(depth_channels, f),
            Format::Custom {
                channels,
                max_value,
            } => write!(
                f,
                "custom format (depth = {channels}, maxval = {max_value})"
            ),
        }
    }
}
impl PartialOrd for Format {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        if let (Format::Standard(lhs), Format::Standard(rhs)) = (self, other) {
            lhs.partial_cmp(rhs)
        } else if self.channels() != other.channels() {
            None
        } else {
            Some(Ord::cmp(&self.max_value(), &other.max_value()))
        }
    }
}

/// Image data.
pub(crate) struct Image<I: Iterator<Item = u16>> {
    /// Size of image.
    size: Size,

    /// Format of image.
    format: Format,

    /// Data for image, with one per channel per pixel.
    data: I,
}
impl<I: Iterator<Item = u16>> Image<I> {
    /// Convert to given depth and channels.
    pub(crate) fn convert(self, format: DepthChannels) -> Image<impl Iterator<Item = u16>> {
        assert!(
            Format::Standard(format) >= self.format,
            "cannot convert {} image into {format} image",
            self.format
        );
        Image {
            size: self.size,
            format: Format::Standard(format),
            data: Convert {
                from: self.format,
                into: format,
                data: self.data,
            }
            .flatten(),
        }
    }

    /// Convert to token stream.
    pub(crate) fn into_tokens(mut self, nesting: Nesting) -> TokenStream {
        let mut data = TokenStream::new();
        for r in 0..self.size.height() {
            if r > 0 {
                data.extend(Some(TokenTree::Punct(Punct::new(',', Spacing::Alone))));
            }
            let mut row = TokenStream::new();
            for c in 0..self.size.width() {
                if c > 0 {
                    row.extend(Some(TokenTree::Punct(Punct::new(',', Spacing::Alone))));
                }
                let mut col = TokenStream::new();
                for idx in 0..(self.format.channels() as usize) {
                    if idx > 0 {
                        col.extend(Some(TokenTree::Punct(Punct::new(',', Spacing::Alone))));
                    }
                    let num = self.data.next().unwrap_or_else(|| {
                        unreachable!("ran out of data at row {r}, col {c}, channel {idx}")
                    });
                    col.extend(Some(match self.format.approx_depth() {
                        Depth::One => TokenTree::Ident(Ident::new(
                            if num == 1 { "true" } else { "false" },
                            Span::call_site(),
                        )),
                        Depth::Eight => TokenTree::Literal(Literal::u8_suffixed(num as u8)),
                        Depth::Sixteen => TokenTree::Literal(Literal::u16_suffixed(num)),
                    }));
                }
                if nesting == Nesting::Inline && self.format.channels() == 1 {
                    row.extend(col);
                } else {
                    row.extend(Some(TokenTree::Group(Group::new(Delimiter::Bracket, col))));
                }
            }
            data.extend(Some(TokenTree::Group(Group::new(Delimiter::Bracket, row))));
        }
        #[cfg(debug_assertions)]
        {
            let remaining = self.data.collect::<Vec<_>>();
            assert_eq!(remaining, []);
        }
        TokenTree::Group(Group::new(Delimiter::Bracket, data)).into()
    }
}

/// Loads size from reader.
fn load_pbm_header(reader: &mut Reader) -> Size {
    let width = reader.read_token(u16::MAX);
    let height = reader.read_token(u16::MAX);
    Size::new(width, height).unwrap_or_else(|| panic!("{reader} was zero-size image"))
}

/// Loads size and maxval from reader.
fn load_pgm_header(reader: &mut Reader) -> (Size, Format) {
    let size = load_pbm_header(reader);
    let Some(depth) = NonZero::new(reader.read_token(u16::MAX)) else {
        panic!("{reader} had maxval = 0");
    };
    let format = Format::new(const { NonZero::new(1).unwrap() }, depth);
    (size, format)
}

/// Loads size and maxval from reader.
fn load_ppm_header(reader: &mut Reader) -> (Size, Format) {
    let size = load_pbm_header(reader);
    let Some(depth) = NonZero::new(reader.read_token(u16::MAX)) else {
        panic!("{reader} had maxval = 0")
    };
    let format = Format::new(const { NonZero::new(3).unwrap() }, depth);
    (size, format)
}

/// Loads PAM header from reader.
fn load_pam_header(reader: &mut Reader) -> (Size, Format) {
    let mut width = None;
    let mut height = None;
    let mut depth = None;
    let mut maxval = None;
    loop {
        match reader.read_token(Header) {
            HeaderDirective::EndHdr => {
                let width =
                    width.unwrap_or_else(|| panic!("{reader} did not have WIDTH specified"));
                let height =
                    height.unwrap_or_else(|| panic!("{reader} did not have HEIGHT specified"));
                let depth =
                    depth.unwrap_or_else(|| panic!("{reader} did not have DEPTH specified"));
                let maxval =
                    maxval.unwrap_or_else(|| panic!("{reader} did not have MAXVAL specified"));
                let format = Format::new(depth, maxval);
                return (
                    Size::new(width, height).unwrap_or_else(|| panic!("{reader} had zero size")),
                    format,
                );
            }
            HeaderDirective::Width => {
                assert!(width.is_none(), "{reader} specified WIDTH twice");
                width = Some(reader.read_token(u16::MAX));
            }
            HeaderDirective::Height => {
                assert!(height.is_none(), "{reader} specified HEIGHT twice");
                height = Some(reader.read_token(u16::MAX));
            }
            HeaderDirective::Depth => {
                assert!(depth.is_none(), "{reader} specified DEPTH twice");
                depth = Some(
                    NonZero::new(reader.read_token(u8::MAX as u16) as u8)
                        .unwrap_or_else(|| panic!("{reader} had DEPTH = 0")),
                );
            }
            HeaderDirective::MaxVal => {
                assert!(maxval.is_none(), "{reader} specified MAXVAL twice");
                maxval = Some(
                    NonZero::new(reader.read_token(u16::MAX))
                        .unwrap_or_else(|| panic!("{reader} had MAXVAL = 0")),
                );
            }
            HeaderDirective::TuplType => {
                reader.read_token(TuplType);
            }
        }
    }
}

/// Error reading bits from PBM.
struct BitError(u8);
impl fmt::Display for BitError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "expected bit, found '{}'", self.0.escape_ascii())
    }
}

/// Load plain PBM file after magic.
pub(crate) fn load_plain_pbm<'a>(reader: &'a mut Reader) -> Image<impl 'a + Iterator<Item = u16>> {
    let size = load_pbm_header(reader);
    let path = reader.path();
    Image {
        size,
        format: Format::Standard(DepthChannels::BlackAndWhite),
        data: reader
            .read_bytes()
            .filter_map(move |b| {
                if b.is_ascii_whitespace() {
                    None
                } else if b == b'0' {
                    // in PBM, 0 means white
                    Some(1)
                } else if b == b'1' {
                    // in PBM, 1 means black
                    Some(0)
                } else {
                    cold::io_failure("parse", &path, &BitError(b))
                }
            })
            .take(size.pixels()),
    }
}

/// Load plain PGM file after magic.
pub(crate) fn load_plain_pgm<'a>(reader: &'a mut Reader) -> Image<impl 'a + Iterator<Item = u16>> {
    let (size, format) = load_pgm_header(reader);
    Image {
        size,
        format,
        data: reader.read_tokens(format.max_value()).take(size.pixels()),
    }
}

/// Load plain PPM file after magic.
pub(crate) fn load_plain_ppm<'a>(reader: &'a mut Reader) -> Image<impl 'a + Iterator<Item = u16>> {
    let (size, format) = load_ppm_header(reader);
    Image {
        size,
        format,
        data: reader.read_tokens(format.max_value()).take(size.values(3)),
    }
}

/// Get bits from byte.
struct PnmByteBits {
    /// Byte being read.
    ///
    /// Constantly shifted right so that the next bit to read is the most significant.
    byte: u8,

    /// Number of significant bits left to read.
    bits_left: u8,
}
impl Iterator for PnmByteBits {
    type Item = bool;
    fn next(&mut self) -> Option<Self::Item> {
        if self.bits_left > 0 {
            let ret = self.byte & 1 == 0;
            self.byte >>= 1;
            self.bits_left -= 1;
            Some(ret)
        } else {
            None
        }
    }
}

/// Load PBM file after magic.
pub(crate) fn load_pbm<'a>(reader: &'a mut Reader) -> Image<impl 'a + Iterator<Item = u16>> {
    let size = load_pbm_header(reader);
    let width_bits = size.width().next_multiple_of(8);
    let width_bytes = width_bits / 8;
    let overflow = if width_bits == size.width() {
        0
    } else {
        (width_bits - size.width()) as u8
    };
    Image {
        size,
        format: Format::Standard(DepthChannels::BlackAndWhite),
        data: reader
            .read_bytes()
            .take(width_bytes * size.height())
            .enumerate()
            .flat_map(move |(idx, byte)| PnmByteBits {
                byte: byte.reverse_bits(),
                bits_left: if (idx + 1) % width_bytes == 0 {
                    8 - overflow
                } else {
                    8
                },
            })
            .map(|b| b as u16),
    }
}

/// Read one or two bytes depending on depth.
struct VariableBytes<I> {
    /// Bit depth.
    depth: Depth,

    /// Iterator over bytes.
    iter: I,
}
impl<I: Iterator<Item = u8>> Iterator for VariableBytes<I> {
    type Item = u16;
    fn next(&mut self) -> Option<Self::Item> {
        if self.depth == Depth::Sixteen {
            let first = self.iter.next()?;
            let second = self.iter.next()?;
            Some(u16::from_be_bytes([first, second]))
        } else {
            Some(self.iter.next()? as u16)
        }
    }
}

/// Load PGM file after magic.
pub(crate) fn load_pgm<'a>(reader: &'a mut Reader) -> Image<impl 'a + Iterator<Item = u16>> {
    let (size, format) = load_pgm_header(reader);
    let depth = format.approx_depth();
    let bytes_per = 1 + (depth == Depth::Sixteen) as usize;
    let iter = reader.read_bytes().take(size.values(1) * bytes_per);
    Image {
        size,
        format,
        data: VariableBytes { depth, iter },
    }
}

/// Load PPM file after magic.
pub(crate) fn load_ppm<'a>(reader: &'a mut Reader) -> Image<impl 'a + Iterator<Item = u16>> {
    let (size, format) = load_ppm_header(reader);
    let depth = format.approx_depth();
    let bytes_per = 1 + (depth == Depth::Sixteen) as usize;
    let iter = reader.read_bytes().take(size.values(3) * bytes_per);
    Image {
        size,
        format,
        data: VariableBytes { depth, iter },
    }
}

/// Load PAM file after magic.
pub(crate) fn load_pam<'a>(reader: &'a mut Reader) -> Image<impl 'a + Iterator<Item = u16>> {
    let (size, format) = load_pam_header(reader);
    let depth = format.approx_depth();
    let bytes_per = 1 + (depth == Depth::Sixteen) as usize;
    let iter = reader
        .read_bytes()
        .take(size.values(format.channels()) * bytes_per);
    Image {
        size,
        format,
        data: VariableBytes { depth, iter },
    }
}