use std::{fmt, fmt::Display, mem::transmute};

mod strategies;
pub use strategies::FilterStrategy;

/// PNG delta filters
#[repr(u8)]
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Hash)]
pub enum RowFilter {
    None,
    Sub,
    Up,
    Average,
    Paeth,
}

impl TryFrom<u8> for RowFilter {
    type Error = ();

    fn try_from(value: u8) -> Result<Self, Self::Error> {
        if value > 4 {
            return Err(());
        }
        unsafe { transmute(value as i8) }
    }
}

impl Display for RowFilter {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        Display::fmt(
            match self {
                Self::None => "None",
                Self::Sub => "Sub",
                Self::Up => "Up",
                Self::Average => "Average",
                Self::Paeth => "Paeth",
            },
            f,
        )
    }
}

impl RowFilter {
    pub(crate) const ALL: [Self; 5] = [Self::None, Self::Sub, Self::Up, Self::Average, Self::Paeth];

    pub(crate) fn filter_line(
        self,
        bpp: usize,
        data: &mut [u8],
        prev_line: &[u8],
        buf: &mut Vec<u8>,
        alpha_bytes: usize,
    ) {
        assert!(data.len() >= bpp);
        assert_eq!(data.len(), prev_line.len());

        if alpha_bytes != 0 {
            self.optimize_alpha(bpp, data, prev_line, bpp - alpha_bytes);
        }

        buf.reserve(data.len() + 1);
        buf.push(self as u8);
        match self {
            Self::None => {
                buf.extend_from_slice(data);
            }
            Self::Sub => {
                buf.extend_from_slice(&data[0..bpp]);
                buf.extend(
                    data.iter()
                        .skip(bpp)
                        .zip(data.iter())
                        .map(|(cur, last)| cur.wrapping_sub(*last)),
                );
            }
            Self::Up => {
                buf.extend(
                    data.iter()
                        .zip(prev_line.iter())
                        .map(|(cur, last)| cur.wrapping_sub(*last)),
                );
            }
            Self::Average => {
                buf.extend(
                    data.iter()
                        .zip(prev_line.iter())
                        .take(bpp)
                        .map(|(cur, &up)| cur.wrapping_sub(up >> 1)),
                );
                buf.extend(data.iter().enumerate().skip(bpp).map(|(i, &cur)| {
                    let left = data[i - bpp];
                    let up = prev_line[i];
                    cur.wrapping_sub(((u16::from(left) + u16::from(up)) >> 1) as u8)
                }));
            }
            Self::Paeth => {
                buf.extend(
                    data.iter()
                        .zip(prev_line.iter())
                        .take(bpp)
                        .map(|(cur, &up)| cur.wrapping_sub(up)),
                );
                buf.extend(data.iter().enumerate().skip(bpp).map(|(i, &cur)| {
                    let left = data[i - bpp];
                    let up = prev_line[i];
                    let left_up = prev_line[i - bpp];
                    cur.wrapping_sub(paeth_predictor(left, up, left_up))
                }));
            }
        }
    }

    // Optimize fully transparent pixels of a scanline such that they will be zeroed when filtered
    fn optimize_alpha(self, bpp: usize, data: &mut [u8], prev_line: &[u8], color_bytes: usize) {
        if self == Self::None {
            // Assume transparent pixels already set to 0
            return;
        }

        let mut pixels: Vec<_> = data.chunks_exact_mut(bpp).collect();
        let prev_pixels: Vec<_> = prev_line.chunks_exact(bpp).collect();
        for i in 0..pixels.len() {
            if pixels[i].iter().skip(color_bytes).all(|b| *b == 0) {
                // If the first pixel in the row is transparent, find the next non-transparent pixel and pretend
                // it is the previous one. This can help improve effectiveness of the Sub and Paeth filters.
                let prev = match i {
                    0 => pixels
                        .iter()
                        .position(|px| px.iter().skip(color_bytes).any(|b| *b != 0))
                        .unwrap_or(i),
                    _ => i - 1,
                };

                // These assertions help eliminate a few bounds checks in the slice accesses below
                assert!(prev < pixels.len());
                assert!(i < prev_pixels.len());

                match self {
                    Self::None => unreachable!(),
                    Self::Sub => {
                        // The code below is roughly equivalent to pixels[i][0..color_bytes].copy_from_slice(&pixels[prev][0..color_bytes]),
                        // if such a thing was possible to do without violating Rust aliasing rules. See:
                        // https://users.rust-lang.org/t/problem-borrowing-two-elements-of-vec-mutably/21446/2

                        if prev < i {
                            let (pixels_head, pixels_tail) = pixels.split_at_mut(prev + 1);
                            pixels_tail[i - prev - 1][0..color_bytes]
                                .copy_from_slice(&pixels_head[prev][0..color_bytes]);
                        } else if prev > i {
                            let (pixels_head, pixels_tail) = pixels.split_at_mut(i + 1);
                            pixels_head[i][0..color_bytes]
                                .copy_from_slice(&pixels_tail[prev - i - 1][0..color_bytes]);
                        } else {
                            // If prev == i, we'd be copying the pixels onto themselves, which is useless
                        }
                    }
                    Self::Up => {
                        pixels[i][0..color_bytes].copy_from_slice(&prev_pixels[i][0..color_bytes]);
                    }
                    Self::Average => {
                        for j in 0..color_bytes {
                            pixels[i][j] = match i {
                                0 => prev_pixels[i][j] >> 1,
                                _ => {
                                    ((u16::from(pixels[i - 1][j]) + u16::from(prev_pixels[i][j]))
                                        >> 1) as u8
                                }
                            };
                        }
                    }
                    Self::Paeth => {
                        for j in 0..color_bytes {
                            pixels[i][j] = match i {
                                0 => pixels[prev][j].min(prev_pixels[i][j]),
                                _ => paeth_predictor(
                                    pixels[i - 1][j],
                                    prev_pixels[i][j],
                                    prev_pixels[i - 1][j],
                                ),
                            };
                        }
                    }
                }
            }
        }
    }

    pub(crate) fn unfilter_line(
        self,
        bpp: usize,
        data: &[u8],
        prev_line: &[u8],
        buf: &mut Vec<u8>,
    ) {
        assert!(data.len() >= bpp);
        assert_eq!(data.len(), prev_line.len());
        let offset = buf.len();
        buf.reserve(data.len());
        match self {
            Self::None => {
                buf.extend_from_slice(data);
            }
            Self::Sub => {
                buf.extend_from_slice(&data[0..bpp]);
                for i in bpp..data.len() {
                    let left = buf[offset + i - bpp];
                    buf.push(data[i].wrapping_add(left));
                }
            }
            Self::Up => {
                buf.extend(
                    data.iter()
                        .zip(prev_line)
                        .map(|(&cur, &last)| cur.wrapping_add(last)),
                );
            }
            Self::Average => {
                for (cur, &up) in data.iter().take(bpp).zip(prev_line.iter().take(bpp)) {
                    buf.push(cur.wrapping_add(up >> 1));
                }
                for i in bpp..data.len() {
                    let left = buf[offset + i - bpp];
                    let up = prev_line[i];
                    buf.push(data[i].wrapping_add(((u16::from(left) + u16::from(up)) >> 1) as u8));
                }
            }
            Self::Paeth => {
                for (cur, &up) in data.iter().take(bpp).zip(prev_line.iter().take(bpp)) {
                    buf.push(cur.wrapping_add(up));
                }
                for i in bpp..data.len() {
                    let left = buf[offset + i - bpp];
                    let up = prev_line[i];
                    let left_up = prev_line[i - bpp];
                    buf.push(data[i].wrapping_add(paeth_predictor(left, up, left_up)));
                }
            }
        }
    }
}

fn paeth_predictor(a: u8, b: u8, c: u8) -> u8 {
    let p = i32::from(a) + i32::from(b) - i32::from(c);
    let pa = (p - i32::from(a)).abs();
    let pb = (p - i32::from(b)).abs();
    let pc = (p - i32::from(c)).abs();
    if pa <= pb && pa <= pc {
        a
    } else if pb <= pc {
        b
    } else {
        c
    }
}
