fn main() {
    println!("Hello, world!");
}

#[derive(Clone, Copy, Default, Debug, PartialEq)]
struct PixelRGBA {
    r: u8,
    g: u8,
    b: u8,
    a: u8,
}

impl PixelRGBA {
    fn new(r: u8, g: u8, b: u8, a: u8) -> Self {
        Self { r, g, b, a }
    }

    fn zero() -> Self {
        Self::new(0, 0, 0, 0)
    }
}

const QOI_OP_RGB: u8 = 0b1111_1110;
const QOI_OP_RGBA: u8 = 0b1111_1111;
const QOI_OP_INDEX: u8 = 0b0000_0000;
const QOI_OP_DIFF: u8 = 0b0100_0000;
const QOI_OP_LUMA: u8 = 0b1000_0000;
const QOI_OP_RUN: u8 = 0b1100_0000;
const QOI_OP_SMALL_MASK: u8 = 0b1100_0000; // mask for the small op codes

struct Decoder<I: Iterator<Item = u8>> {
    // QOI codec state information
    back_buffer: [PixelRGBA; 64],
    prev_pixel: PixelRGBA,

    bytes: I,
    run_len: u8,
}

impl<I> Decoder<I>
where
    I: Iterator<Item = u8>,
{
    fn new(bytes: I) -> Self {
        Self {
            back_buffer: [PixelRGBA::zero(); 64],
            prev_pixel: PixelRGBA {
                r: 0,
                g: 0,
                b: 0,
                a: 255,
            },
            bytes,
            run_len: 0,
        }
    }

    // A hack to unit test the index lookup behavior. A partial test can be done
    // to verify the basic indexing principles by preloading a known buffer and
    // then extracting back-referenced data out of it. A complete test should
    // feed in other valid operations that populate the backbuffer, and then index
    // op codes to demonstrate the indexing operations.
    fn new_with_backbuffer(bytes: I, back_buffer: [PixelRGBA; 64]) -> Self {
        Self {
            back_buffer,
            prev_pixel: PixelRGBA {
                r: 0,
                g: 0,
                b: 0,
                a: 255,
            },
            bytes,
            run_len: 0,
        }
    }

    // A hack to unit test the run behavior. Same idea as the new_with_backbuffer()
    // function, but for testing a run of pixels.
    fn new_with_previous_pixel(bytes: I, prev_pixel: PixelRGBA) -> Self {
        Self {
            back_buffer: [PixelRGBA::zero(); 64],
            prev_pixel,
            bytes,
            run_len: 0,
        }
    }

    pub(crate) fn peek_prev_pixel(&self) -> &PixelRGBA {
        &self.prev_pixel
    }

    pub(crate) fn peek_backbuffer(&self, idx: usize) -> &PixelRGBA {
        &self.back_buffer[idx]
    }
}

mod codec_utils {
    use super::PixelRGBA;
    pub(crate) fn hash(pixel: PixelRGBA) -> u8 {
        pixel
            .r
            .wrapping_mul(3)
            .wrapping_add(pixel.g.wrapping_mul(5))
            .wrapping_add(pixel.b.wrapping_mul(7))
            .wrapping_add(pixel.a.wrapping_mul(11))
            % 64
    }
}

impl<'input, I> Iterator for Decoder<I>
where
    I: Iterator<Item = u8>,
{
    type Item = PixelRGBA;

    fn next(&mut self) -> Option<Self::Item> {
        if self.run_len > 0 {
            self.run_len -= 1;
            return Some(self.prev_pixel);
        } else {
            // Two kinds of patterns to match:
            // 1. Whole byte tag -- RGB and RGBA
            // 2. Partial byte tag -- 2 front bits of Index, Diff, Luma, and Run
            let byte = self.bytes.next()?;
            if byte == QOI_OP_RGB {
                return Some(PixelRGBA {
                    r: self.bytes.next()?,
                    g: self.bytes.next()?,
                    b: self.bytes.next()?,
                    a: self.prev_pixel.a,
                });
            } else if byte == QOI_OP_RGBA {
                return Some(PixelRGBA {
                    r: self.bytes.next()?,
                    g: self.bytes.next()?,
                    b: self.bytes.next()?,
                    a: self.bytes.next()?,
                });
            } else {
                match byte & QOI_OP_SMALL_MASK {
                    QOI_OP_INDEX => {
                        let idx = (byte & !QOI_OP_SMALL_MASK) as usize;
                        return Some(self.back_buffer[idx]);
                    }
                    QOI_OP_DIFF => {
                        let dr = ((byte & 0b0011_0000) >> 4).wrapping_sub(2);
                        let dg = ((byte & 0b0000_1100) >> 2).wrapping_sub(2);
                        let db = (byte & 0b0000_0011).wrapping_sub(2);
                        let result = PixelRGBA {
                            r: self.prev_pixel.r.wrapping_add(dr),
                            g: self.prev_pixel.g.wrapping_add(dg),
                            b: self.prev_pixel.b.wrapping_add(db),
                            a: self.prev_pixel.a,
                        };
                        self.prev_pixel = result;
                        return Some(result);
                    }
                    QOI_OP_LUMA => {
                        let dg = (byte & !QOI_OP_SMALL_MASK).wrapping_sub(32);
                        let packed = self.bytes.next()?;
                        let drdg = ((packed & 0b1111_0000) >> 4).wrapping_sub(8);
                        let dbdg = (packed & 0b0000_1111).wrapping_sub(8);
                        let dr = drdg.wrapping_add(dg);
                        let db = dbdg.wrapping_add(dg);
                        let result = PixelRGBA {
                            r: self.prev_pixel.r.wrapping_add(dr),
                            g: self.prev_pixel.g.wrapping_add(dg),
                            b: self.prev_pixel.b.wrapping_add(db),
                            a: self.prev_pixel.a,
                        };
                        self.prev_pixel = result;
                        return Some(result);
                    }
                    QOI_OP_RUN => {
                        self.run_len = byte & !QOI_OP_SMALL_MASK;
                        // storage bias of -1, so a +1 should be on the end here.
                        // However, I'm immediately popping off the first occurrence
                        // and returning a PixelRGBA, so the count is also immediatly
                        // dropped by 1
                        return Some(self.prev_pixel);
                    }
                    _ => panic!("bad op code{}", byte),
                }
            }
        }
    }
}

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

    #[test] // this is mostly just to drive the function. Make sure it wraps or crashes in debug.
    fn test_backref_hash_function() {
        let pixel = PixelRGBA {
            r: 100,
            g: 80,
            b: 90,
            a: 255,
        };
        let expected = 39;
        assert_eq!(codec_utils::hash(pixel), expected);
    }

    #[test]
    fn decoder_unpack_rgb() {
        // compressed RGB values should be expanded back out to RGBA
        // with an assumed alpha of 0xFF.
        let compressed = [
            QOI_OP_RGB, 0xFF, 0xFF, 0xFF, QOI_OP_RGB, 0x7F, 0x00, 0xAD, QOI_OP_RGB, 0x00, 0x00,
            0x00,
        ];

        let expected = [
            PixelRGBA {
                r: 0xFF,
                g: 0xFF,
                b: 0xFF,
                a: 0xFF,
            },
            PixelRGBA {
                r: 0x7F,
                g: 0x00,
                b: 0xAD,
                a: 0xFF,
            },
            PixelRGBA {
                r: 0x00,
                g: 0x00,
                b: 0x00,
                a: 0xFF,
            },
        ];

        let decoder = Decoder::new(compressed.into_iter());
        let result = decoder.collect::<Vec<PixelRGBA>>();
        assert_eq!(result, expected);
    }

    #[test]
    fn decoder_unpack_rgba() {
        let compressed = [
            QOI_OP_RGBA,
            0xFF,
            0xFF,
            0xFF,
            0xFF,
            QOI_OP_RGBA,
            0x7F,
            0x7F,
            0x7F,
            0xFF,
            QOI_OP_RGBA,
            0x10,
            0x20,
            0x30,
            0x40,
        ];

        let expected = [
            PixelRGBA {
                r: 0xFF,
                g: 0xFF,
                b: 0xFF,
                a: 0xFF,
            },
            PixelRGBA {
                r: 0x7f,
                g: 0x7f,
                b: 0x7f,
                a: 0xFF,
            },
            PixelRGBA {
                r: 0x10,
                g: 0x20,
                b: 0x30,
                a: 0x40,
            },
        ];

        let decoder = Decoder::new(compressed.into_iter());
        let result: Vec<PixelRGBA> = decoder.collect();
        assert_eq!(result, expected);
    }

    #[test]
    fn decoder_unpack_index() {
        let mut backbuffer = [PixelRGBA::zero(); 64];
        backbuffer[0] = PixelRGBA::new(255, 255, 255, 255);
        backbuffer[1] = PixelRGBA::new(0, 255, 0, 255);
        backbuffer[2] = PixelRGBA::new(255, 0, 255, 255);
        backbuffer[3] = PixelRGBA::new(0, 0, 0, 0);
        backbuffer[10] = PixelRGBA::new(10, 10, 10, 0);
        backbuffer[11] = PixelRGBA::new(0xF0, 0x2C, 0xAF, 0xFF);

        let compressed = [
            (QOI_OP_INDEX | 0),
            (QOI_OP_INDEX | 11),
            (QOI_OP_INDEX | 1),
            (QOI_OP_INDEX | 2),
            (QOI_OP_INDEX | 3),
            (QOI_OP_INDEX | 10),
            (QOI_OP_INDEX | 42),
        ];

        let expected = [
            PixelRGBA::new(255, 255, 255, 255),
            PixelRGBA::new(0xF0, 0x2C, 0xAF, 0xFF),
            PixelRGBA::new(0, 0xFF, 0, 0xFF),
            PixelRGBA::new(255, 0, 255, 255),
            PixelRGBA::new(0, 0, 0, 0),
            PixelRGBA::new(10, 10, 10, 0),
            PixelRGBA::zero(),
        ];

        let decoder = Decoder::new_with_backbuffer(compressed.into_iter(), backbuffer);
        let result: Vec<PixelRGBA> = decoder.collect();
        assert_eq!(result, expected);
    }

    #[test]
    fn decoder_unpack_diff() {
        // DIFF components are 2 bit values with a bias of 2.
        // i.e. : 0b00 is -2, and 0b11 is +1
        let compressed = [
            (QOI_OP_DIFF | 0b0011_1111), // (1, 1, 1)
            (QOI_OP_DIFF | 0b0011_1010), // (1, 0, 0)
            (QOI_OP_DIFF | 0b0010_1110), // (0, 1, 0)
            (QOI_OP_DIFF | 0b0010_1011), // (0, 0, 1)
            (QOI_OP_DIFF | 0b0011_1011), // (1, 0, 1)
            (QOI_OP_DIFF | 0b0001_1001), // (-1, 0, -1)
            (QOI_OP_DIFF | 0b0000_0000), // (-2, -2, -2)
        ];

        // the codec begins with a pixel at (0, 0, 0, 255), so these results
        // are diffs from that.
        let expected = [
            PixelRGBA::new(1, 1, 1, 255),
            PixelRGBA::new(2, 1, 1, 255),
            PixelRGBA::new(2, 2, 1, 255),
            PixelRGBA::new(2, 2, 2, 255),
            PixelRGBA::new(3, 2, 3, 255),
            PixelRGBA::new(2, 2, 2, 255),
            PixelRGBA::new(0, 0, 0, 255),
        ];

        let decoder = Decoder::new(compressed.into_iter());
        let result: Vec<PixelRGBA> = decoder.collect();
        assert_eq!(result, expected);
    }

    #[test]
    fn decoder_unpack_diff_rollover() {
        let init_pixel = PixelRGBA::new(255, 255, 255, 255);
        let compressed = [
            (QOI_OP_DIFF | 0b0011_1111), // +1s
            (QOI_OP_DIFF | 0b0011_1111), // +1s
            (QOI_OP_DIFF | 0b0010_1010), // +0s, could have been an index or a run, probably
        ];
        let expected = [
            PixelRGBA::new(0, 0, 0, 255), // 255 rollover to 0
            PixelRGBA::new(1, 1, 1, 255), // +1s
            PixelRGBA::new(1, 1, 1, 255), // holds at 1s
        ];

        let decoder = Decoder::new_with_previous_pixel(compressed.into_iter(), init_pixel);
        let result: Vec<PixelRGBA> = decoder.collect();
        assert_eq!(result, expected);
    }

    #[test]
    fn decoder_unpack_diff_rollunder() {
        let init_pixel = PixelRGBA::new(0, 0, 0, 255);
        let compressed = [
            (QOI_OP_DIFF | 0b0001_0101), // -1s
            (QOI_OP_DIFF | 0b0001_0101), // -1s
            (QOI_OP_DIFF | 0b0010_1010), // 0s
        ];
        let expected = [
            PixelRGBA::new(255, 255, 255, 255),
            PixelRGBA::new(254, 254, 254, 255),
            PixelRGBA::new(254, 254, 254, 255),
        ];

        let decoder = Decoder::new_with_previous_pixel(compressed.into_iter(), init_pixel);
        let result: Vec<PixelRGBA> = decoder.collect();
        assert_eq!(result, expected);
    }

    #[test]
    fn decoder_unpack_luma() {
        // red and blue diffs are relative to the green channel as (dr - dg) and (db - dg)
        // Their finished diffs need to invert this operation.
        // Diff(dg, dr-dg, db-dg) and Pix (dr, dg, db)
        let compressed = [
            (QOI_OP_LUMA | 0b0011_1111),
            (0b1111_1111), // Diff( 31, 7,  7) -> Pix (38, 31, 38)
            (QOI_OP_LUMA | 0b0010_0000),
            (0b1000_1000), // Diff(  0, 0,  0) -> Pix (0, 0, 0)
            (QOI_OP_LUMA | 0b0010_0001),
            (0b1111_1111), // Diff(  1, 7,  7) -> Pix (8, 1, 8)
            (QOI_OP_LUMA | 0b0001_0011),
            (0b1100_0011), // Diff(-13, 4, -5) -> Pix (-9, -13, -18)
        ];

        let expected = [
            PixelRGBA::new(38, 31, 38, 255),
            PixelRGBA::new(38, 31, 38, 255),
            PixelRGBA::new(46, 32, 46, 255),
            PixelRGBA::new(37, 19, 28, 255),
        ];

        let decoder = Decoder::new(compressed.into_iter());
        let result: Vec<PixelRGBA> = decoder.collect();
        assert_eq!(result, expected);
    }

    #[test]
    fn decoder_unpack_luma_rollover() {
        let init_pixel = PixelRGBA::new(255, 255, 255, 255);
        let compressed = [
            (QOI_OP_LUMA | 0b0011_1111),
            (0b1111_1111), // Diff (31, 7, 7) -> Pix (38, 31, 38)
        ];
        let expected = PixelRGBA::new(37, 30, 37, 255);

        let mut decoder = Decoder::new_with_previous_pixel(compressed.into_iter(), init_pixel);
        let result = decoder
            .next()
            .expect("Oops, didn't get a Pixel back from the Decoder");
        assert_eq!(result, expected);
    }

    #[test]
    fn decoder_unpack_luma_rollunder() {
        let init_pixel = PixelRGBA::new(0, 0, 0, 255);
        let compressed = [
            (QOI_OP_LUMA | 0b0001_0011),
            (0b1100_0011), // Diff(-13, 4, -5) -> Pix (-9, -13, -18)
        ];
        let expected = PixelRGBA::new(247, 243, 238, 255);

        let mut decoder = Decoder::new_with_previous_pixel(compressed.into_iter(), init_pixel);
        let result = decoder
            .next()
            .expect("Oops, didn't get a Pixel back from the Decoder");
        assert_eq!(result, expected);
    }

    #[test]
    fn decoder_unpack_run() {
        let compressed = [
            (QOI_OP_RUN | 0b0000_0000), //  1 -- bias of -1, so all zeros is a run of 1 pixel
            (QOI_OP_RUN | 0b0000_1100), // 13 // 0b1111?  no it isn't. what?
        ];

        let init_pixel = PixelRGBA::new(50, 100, 150, 200);
        let expected = [init_pixel; 14];

        // the run instructions should really have been collapsed into just one, but
        // lets pretend an encoder did this for some reason. The decoder can still
        // unpack this correctly, it's just a sub-optimal compression is all.

        let decoder = Decoder::new_with_previous_pixel(compressed.into_iter(), init_pixel);
        let result: Vec<PixelRGBA> = decoder.collect();
        assert_eq!(result, expected);
    }

    #[test]
    fn decoder_prev_pixel_verify() {
        let compressed = [
            QOI_OP_RGB, 0x10, 0x10, 0x10,
            QOI_OP_RGBA, 0x20, 0x20, 0x20, 0x20,
            (QOI_OP_INDEX | 1),
            (QOI_OP_DIFF | 0b0011_1111),
            (QOI_OP_LUMA | 0b0011_111), 0b0011_1111,
            (QOI_OP_RUN | 2),
            QOI_OP_RGBA, 0xFF, 0xFF, 0xFF, 0xFF,
        ];

        let expected = [
            PixelRGBA::new(0, 0, 0, 0xFF),          // init
            PixelRGBA::new(0x10, 0x10, 0x10, 0xFF), // RGB 
            PixelRGBA::new(0x20, 0x20, 0x20, 0xFF), // RGBA
            PixelRGBA::new(0, 0, 0, 0),             // INDEX -- this doubles as a small test for the backbuffer operation
            PixelRGBA::new(0x1, 0x1, 0x1, 0x0),     // DIFF
            PixelRGBA::new(0x27, 0x20, 0x27, 0x0),   // LUMA
            PixelRGBA::new(0x27, 0x20, 0x27, 0x0),   // RUN 1
            PixelRGBA::new(0x27, 0x20, 0x27, 0x0),   // RUN 2
            PixelRGBA::new(0x27, 0x20, 0x27, 0x0),   // RUN 3
            // final OP_RGBA is just to flush out the OP_RUN prev_pixel value
        ];

        let mut decoder = Decoder::new(compressed.into_iter());

        let mut result = Vec::<PixelRGBA>::new();
        loop {
            if let Some(_) = decoder.next() {
                result.push(*decoder.peek_prev_pixel());
            } else {
                break;
            }
        }
        assert_eq!(result, expected);
    }

    #[test]
    fn decoder_backbuffer_verify() {
        let compressed = [
            QOI_OP_RGB, 0x10, 0x10, 0x10,
            QOI_OP_RGBA, 0x20, 0x20, 0x20, 0x20,
            (QOI_OP_INDEX | 1),
            (QOI_OP_DIFF | 0b0011_1111),
            (QOI_OP_LUMA | 0b0011_111), 0b0011_1111,
            (QOI_OP_RUN | 2),
            QOI_OP_RGBA, 0xFF, 0xFF, 0xFF, 0xFF,
        ];
        // these are the indices where we're expecting each pixel to land.
        // Each pixel gets put into this backbuffer as it's en/de-coded.
        // For RGB and RGBA, it'll simply assign a value into the index.
        // For INDEX, the write can be skipped, and the expected index will be
        //  the same as the one in the op code. By definition, it has to be.
        // OP_DIFF & OP_LUMA need to consider the value in the backbuffer, as
        //  they'll be using it to compute the new pixel.
        //  
        let indices = [
            31, // Pix (16, 16, 16, 255)
            0,  // Pix (32, 32, 32, 32)
            1,  // Pix (0, 0, 0, 0)
            3,  // current state: Pix (1, 1, 1, 0)
            38, // current state: Pix(39, 32, 39, 0)
            38, // run x1 
            38, // run x2
            38, // run x3
        ];
        let mut decoder = Decoder::new(compressed.into_iter());
        let mut iters = 0;
        loop {
            if let Some(pixel) = decoder.next() {
                // pixel has been decompressed, so it should be in the backbuffer by this point

                // query it out:
                let stored_px = decoder.peek_backbuffer(indices[iters]);
                
                // and compare it to the value returned from iteration
                assert_eq!(&pixel, stored_px);
            } else {
                break;
            }
            iters += 1;
        }
    }
}