use std::iter::FusedIterator;

use image::{Rgb, RgbImage};
use thiserror::Error;

use crate::Layout;

#[derive(Clone, Copy, Debug, PartialEq, Eq, Error)]
#[allow(clippy::module_name_repetitions)]
pub enum StrandifierError {
    #[error("Image too large")]
    ImageTooLarge,
    #[error("Invalid strand number")]
    InvalidStrand,
}

/// Iterator over a strand's pixels, extracted from an [`RgbImage`].
#[derive(Clone, Debug)]
pub struct Strandifier<'a> {
    layout: Layout,
    image: &'a RgbImage,

    pixels_remaining: u32,
    next_x: u32,
    next_y: u32,

    offset_x: u32,
    offset_y: u32,
}

impl<'a> Strandifier<'a> {
    /// Create an iterator over a strand's pixels.
    ///
    /// # Errors
    ///
    /// Returns `StrandifierError::WrongDimensions` if the `layout` does not match up with the
    /// `image`'s dimensions.
    ///
    /// Returns `StrandifierError::InvalidStrand` if the supplied strand number is higher than the
    /// number of strands in the [`Layout`].
    #[allow(clippy::missing_panics_doc)]
    pub fn make_strand(
        layout: Layout,
        image: &'a RgbImage,
        strand_num: u32,
    ) -> Result<Self, StrandifierError> {
        if layout.width_px() < image.width() || layout.height_px() < image.height() {
            return Err(StrandifierError::ImageTooLarge);
        }

        if strand_num > layout.num_strands() {
            return Err(StrandifierError::InvalidStrand);
        }

        let panel_x = strand_num % layout.num_panels_h;
        let panel_y = strand_num / layout.num_panels_h;

        let first_x = panel_x * layout.gang_len;
        let first_y = panel_y * layout.num_gangs;

        let offset_x = (layout.width_px() - image.width()) / 2;
        let offset_y = (layout.height_px() - image.height()) / 2;

        Ok(Self {
            layout,
            image,

            pixels_remaining: layout.strand_len(),
            next_x: first_x,
            next_y: first_y,

            offset_x,
            offset_y,
        })
    }
}

impl<'a> Iterator for Strandifier<'a> {
    type Item = Rgb<u8>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.pixels_remaining == 0 {
            return None;
        }

        let x = self.next_x;
        let y = self.next_y;

        if (y % self.layout.num_gangs) % 2 == 0 {
            // Right-moving gang
            if x % self.layout.gang_len == self.layout.gang_len - 1 {
                // End of gang - move down
                self.next_y += 1;
            } else {
                // move right
                self.next_x += 1;
            }
        } else {
            // Left-moving gang
            if x % self.layout.gang_len == 0 {
                // End of gang - move down
                self.next_y += 1;
            } else {
                // move left
                self.next_x -= 1;
            }
        }
        self.pixels_remaining -= 1;

        let opt_px: Option<_> = try {
            let x = x.checked_sub(self.offset_x)?;
            let y = y.checked_sub(self.offset_y)?;

            *self.image.get_pixel_checked(x, y)?
        };

        Some(opt_px.unwrap_or(Rgb([0, 0, 0])))
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let remaining = self.pixels_remaining.try_into().unwrap();

        (remaining, Some(remaining))
    }
}

impl<'a> FusedIterator for Strandifier<'a> {}
impl<'a> ExactSizeIterator for Strandifier<'a> {}

#[cfg(test)]
mod tests {
    use super::*;

    fn test_image(layout: Layout) -> RgbImage {
        let width = layout.width_px();
        let height = layout.height_px();

        RgbImage::from_fn(width, height, |x, y| {
            Rgb([x.try_into().unwrap(), y.try_into().unwrap(), 0])
        })
    }

    #[test]
    fn strandifier_basic() {
        let layout = Layout {
            gang_len: 5,
            num_gangs: 4,

            num_panels_h: 3,
            num_panels_v: 6,
        };
        let image = test_image(layout);

        let strand = Strandifier::make_strand(layout, &image, 0).unwrap();
        assert_eq!(
            strand.collect::<Vec<_>>(),
            [
                Rgb([0, 0, 0]),
                Rgb([1, 0, 0]),
                Rgb([2, 0, 0]),
                Rgb([3, 0, 0]),
                Rgb([4, 0, 0]),
                //
                Rgb([4, 1, 0]),
                Rgb([3, 1, 0]),
                Rgb([2, 1, 0]),
                Rgb([1, 1, 0]),
                Rgb([0, 1, 0]),
                //
                Rgb([0, 2, 0]),
                Rgb([1, 2, 0]),
                Rgb([2, 2, 0]),
                Rgb([3, 2, 0]),
                Rgb([4, 2, 0]),
                //
                Rgb([4, 3, 0]),
                Rgb([3, 3, 0]),
                Rgb([2, 3, 0]),
                Rgb([1, 3, 0]),
                Rgb([0, 3, 0]),
            ]
        );

        let strand = Strandifier::make_strand(layout, &image, 4).unwrap();
        assert_eq!(
            strand.collect::<Vec<_>>(),
            [
                Rgb([5, 4, 0]),
                Rgb([6, 4, 0]),
                Rgb([7, 4, 0]),
                Rgb([8, 4, 0]),
                Rgb([9, 4, 0]),
                //
                Rgb([9, 5, 0]),
                Rgb([8, 5, 0]),
                Rgb([7, 5, 0]),
                Rgb([6, 5, 0]),
                Rgb([5, 5, 0]),
                //
                Rgb([5, 6, 0]),
                Rgb([6, 6, 0]),
                Rgb([7, 6, 0]),
                Rgb([8, 6, 0]),
                Rgb([9, 6, 0]),
                //
                Rgb([9, 7, 0]),
                Rgb([8, 7, 0]),
                Rgb([7, 7, 0]),
                Rgb([6, 7, 0]),
                Rgb([5, 7, 0]),
            ]
        );
    }

    #[test]
    fn strandifier_thin() {
        let layout = Layout {
            gang_len: 1,
            num_gangs: 6,

            num_panels_h: 3,
            num_panels_v: 4,
        };
        let image = test_image(layout);

        let strand = Strandifier::make_strand(layout, &image, 0).unwrap();
        assert_eq!(
            strand.collect::<Vec<_>>(),
            [
                Rgb([0, 0, 0]),
                Rgb([0, 1, 0]),
                Rgb([0, 2, 0]),
                Rgb([0, 3, 0]),
                Rgb([0, 4, 0]),
                Rgb([0, 5, 0]),
            ]
        );

        let layout = Layout {
            gang_len: 6,
            num_gangs: 1,

            num_panels_h: 3,
            num_panels_v: 4,
        };
        let image = test_image(layout);

        let strand = Strandifier::make_strand(layout, &image, 11).unwrap();
        assert_eq!(
            strand.collect::<Vec<_>>(),
            [
                Rgb([12, 3, 0]),
                Rgb([13, 3, 0]),
                Rgb([14, 3, 0]),
                Rgb([15, 3, 0]),
                Rgb([16, 3, 0]),
                Rgb([17, 3, 0]),
            ]
        );
    }

    #[test]
    fn strandifier_errors() {
        let layout = Layout {
            gang_len: 5,
            num_gangs: 4,

            num_panels_h: 3,
            num_panels_v: 6,
        };
        let image = test_image(layout);
        assert_eq!(
            Strandifier::make_strand(layout, &image, 20).unwrap_err(),
            StrandifierError::InvalidStrand
        );

        let layout2 = Layout {
            num_gangs: 3,
            ..layout
        };
        assert_eq!(
            Strandifier::make_strand(layout2, &image, 0).unwrap_err(),
            StrandifierError::ImageTooLarge
        );
    }
}