use crate::{
    constants::*,
    DecodeError,
    PixelRGBA,
};

pub fn try_read_magic<I: Iterator<Item=u8>>(bytes: &mut I) -> Result<(), DecodeError> {
    for i in 0..4 {
        if let Some(letter) = bytes.next() {
            if letter != MAGIC[i] {
                return Err(DecodeError::Magic)
            }
        } else {
            return Err(DecodeError::EarlyIteratorExhaustion)
        }
    }
    return Ok(())
}

fn try_read_u32<I: Iterator<Item=u8>>(bytes: &mut I) -> Result<u32, DecodeError> {
    let mut repacker = || {
        if let Some(next_byte) = bytes.next() {
            return Result::<u32, DecodeError>::Ok(next_byte as u32);
        } else {
            return Err(DecodeError::EarlyIteratorExhaustion);
        }
    };

    Ok(
        (repacker()? << 24)
      + (repacker()? << 16)
      + (repacker()? << 8)
      + (repacker()?)
    )
}

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


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

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;
            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 {
                let result = PixelRGBA {
                    r: self.bytes.next()?,
                    g: self.bytes.next()?,
                    b: self.bytes.next()?,
                    a: self.prev_pixel.a,
                };
                self.prev_pixel = result;
                self.back_buffer[codec_utils::hash(result) as usize] = result;
                return Some(result);
            } else if byte == QOI_OP_RGBA {
                let result = PixelRGBA {
                    r: self.bytes.next()?,
                    g: self.bytes.next()?,
                    b: self.bytes.next()?,
                    a: self.bytes.next()?,
                };
                self.prev_pixel = result;
                self.back_buffer[codec_utils::hash(result) as usize] = result;
                return Some(result);
            } else {
                match byte & QOI_OP_SMALL_MASK {
                    QOI_OP_INDEX => {
                        let idx = (byte & !QOI_OP_SMALL_MASK) as usize;
                        self.prev_pixel = self.back_buffer[idx];
                        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;
                        self.back_buffer[codec_utils::hash(result) as usize] = 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;
                        self.back_buffer[codec_utils::hash(result) as usize] = 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::*;
    impl<I> Decoder<I>
    where
        I: Iterator<Item = u8>,
    {
        // 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,
            }
        }

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

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

    #[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_1111),
            0b1111_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, 0x20), // 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
            PixelRGBA::new(0xFF, 0xFF, 0xFF, 0xFF),
        ];

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

        let mut result = Vec::<PixelRGBA>::new();
        loop {
            result.push(*decoder.peek_prev_pixel());
            if let None = decoder.next() {
                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_RGBA,
            0x03,
            0x01,
            0x01,
            0x04, // filler to populate backbuffer[1] for the next OP_INDEX
            (QOI_OP_INDEX | 1),
            (QOI_OP_DIFF | 0b0011_1111), // + Pix (1, 1, 1, 0)
            (QOI_OP_LUMA | 0b0011_1111),
            0b1111_1111, // + Pix (38, 31, 38)
            (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.
        //  *however* it's possible for an index to refer to a value that shouldn't
        //  hash to that location. The backbuffer initialization will cause this
        //  scenario.
        //
        //  Without an OP_INDEX buffer write, this condition will cause codec corruption.
        //  If the Indexed value should have landed somewhere else, and that other location
        //  has a different value in it, then the next operation to reference that field
        //  will receive the wrong data.

        // 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 = [
            37, // Pix (16, 16, 16, 255)
            0,  // Pix (32, 32, 32, 32)
            1,  // Pix (3, 1, 1, 4)
            1,  // Pix (3, 1, 1, 4)
            16, // Pix (4, 2, 2, 4)
            39, // Pix (42, 33, 40, 4)
            39, // run x1
            39, // run x2
            39, // run x3
            38, // final RGBA
        ];
        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;
        }
    }

    #[test]
    fn test_try_read_u32() {
        let input = [10u8, 200u8, 255u8, 1u8, 254u8, 123u8, 45u8];
        let expected =  180_944_641;
        assert_eq!(
            try_read_u32(&mut input.into_iter()).expect("Couldn't read u32 from byte array"),
            expected
        );

        let too_short = [0u8];
        let expected = Result::<u32, DecodeError>::Err(DecodeError::EarlyIteratorExhaustion);
        assert_eq!(
            try_read_u32(&mut too_short.into_iter()),
            expected
        );
    }
}