use core::{
    arch::asm,
    hint::black_box,
    ptr::addr_of,
    slice::{Chunks, Iter},
};

use atmega_hal::{
    Atmega, Usart,
    pac::USART0,
    port::{PB3, PB5, PD0, PD1, Pin, PinOps, mode},
};
use embedded_hal::delay::DelayNs;
use ufmt::derive::uDebug;

use crate::{
    CoreClock,
    display::{Color, Display, Position, Rgb565, Writeable},
    qoi,
};

// progmem! {
//     /// Static byte stored in progmem!
//     static progmem P_BYTE: u8 = 42;
//     static progmem VERY_LARGE: [u8; 23731] = *include_bytes!("../images/cat.qoi");
// }

// https://qoiformat.org/qoi-specification.pdf

#[unsafe(link_section = ".progmem.data")]
static LARGE_CAT: [u8; 23731] = *include_bytes!("../images/cat.qoi");

pub struct LargeIterator {
    ptr: *const u8,
    length: usize,
    index: usize,
}

impl LargeIterator {
    pub fn from<const N: usize>(addr: *const [u8; N]) -> LargeIterator {
        LargeIterator {
            ptr: addr.cast(),
            length: N,
            index: 0,
        }
    }
}

impl Iterator for LargeIterator {
    type Item = u8;

    fn next(&mut self) -> Option<Self::Item> {
        if self.index >= self.length {
            return None;
        }
        let elem_ptr = self.ptr.wrapping_add(self.index);
        let res: u8;
        unsafe {
            asm!(
                "lpm {}, Z",
                out(reg) res,
                in("Z") elem_ptr,
            )
        }
        self.index += 1;
        Some(res)
    }
}

#[derive(Debug, uDebug)]
pub enum QoiErr {
    InvalidMagicNumber,
    UnexpectedEOF,
}

pub fn draw_image<T: DelayNs, DCPin: PinOps, RSTPin: PinOps>(
    serial: &mut Usart<USART0, Pin<mode::Input, PD0>, Pin<mode::Output, PD1>, CoreClock>,
    buffer: &[u8],
    display: &mut Display<T, DCPin, RSTPin>,
    position: Position,
) -> Result<(), QoiErr> {
    // let mut iter = buffer.iter().map(|b| *b);

    let addr = addr_of!(LARGE_CAT);
    let mut iter = LargeIterator::from(addr);

    if let (Some(113), Some(111), Some(105), Some(102)) =
        (iter.next(), iter.next(), iter.next(), iter.next())
    {
        fn next(iter: &mut dyn Iterator<Item = u8>) -> Result<u8, QoiErr> {
            iter.next().ok_or(QoiErr::UnexpectedEOF)
        }

        let width = u32::from_be_bytes([
            next(&mut iter)?,
            next(&mut iter)?,
            next(&mut iter)?,
            next(&mut iter)?,
        ]) as u16;
        let height = u32::from_be_bytes([
            next(&mut iter)?,
            next(&mut iter)?,
            next(&mut iter)?,
            next(&mut iter)?,
        ]) as u16;
        let _channels = next(&mut iter)?;
        let _colorspace = next(&mut iter)?;

        ufmt::uwriteln!(serial, "Successfully read QOI header").unwrap();

        let scale_factor = 240 / width;

        display.set_window(
            position,
            Position {
                x: position.x + (width * scale_factor) - 1,
                y: position.y + (height * scale_factor) - 1,
            },
        );

        let qoi_iter = QoiIterator::from(&mut iter, width * height);
        let scale_iter = ScaleIterator {
            qoi: qoi_iter,
            last_row: [Rgb565::yellow(); 120],
            scale_factor: scale_factor as usize,
            width: width as usize,
            counter: 0,
            index: 0,
        };

        let mut counter = 0u32;
        for pixel in scale_iter {
            let [c1, c2] = pixel.as_color().bytes;
            display.write(Writeable::Data(&[c1, c2]));
            counter += 1;
        }
        ufmt::uwriteln!(serial, "Counter: {}", counter).unwrap();

        Ok(())
    } else {
        Err(QoiErr::InvalidMagicNumber)
    }
}

struct ScaleIterator<'a> {
    qoi: QoiIterator<'a>,
    last_row: [Rgb565; 120],
    width: usize,
    scale_factor: usize,
    counter: usize,
    index: usize,
}

impl<'a> Iterator for ScaleIterator<'a> {
    type Item = Rgb565;

    fn next(&mut self) -> Option<Self::Item> {
        if self.scale_factor == 1 {
            return self.qoi.next();
        }

        if self.index >= self.width * self.scale_factor {
            self.counter = (self.counter + 1) % self.scale_factor;
            self.index = 0;
        }
        let index_div = self.index / self.scale_factor;

        if self.counter % self.scale_factor == 0 {
            if (self.index % self.scale_factor) == 0 {
                if let Some(pixel) = self.qoi.next() {
                    self.last_row[index_div] = pixel;
                    self.index += 1;
                    Some(pixel)
                } else {
                    None
                }
            } else {
                let pixel = self.last_row[index_div];
                self.index += 1;
                Some(pixel)
            }
        } else {
            let pixel = self.last_row[index_div];
            self.index += 1;
            Some(pixel)
        }
    }
}

struct QoiIterator<'a> {
    inner: &'a mut dyn Iterator<Item = u8>,
    prev_pixels: [Rgb565; 64],
    last_pixel: Rgb565,
    repeat: u8,
    expected_colors: u16,
    parsed_colors: u16,
}

impl<'a> QoiIterator<'a> {
    fn from(bytes: &'a mut dyn Iterator<Item = u8>, expected: u16) -> Self {
        QoiIterator {
            inner: bytes,
            prev_pixels: [Rgb565(0, 255, 0); 64],
            last_pixel: Rgb565(0, 0, 0),
            repeat: 0,
            expected_colors: expected,
            parsed_colors: 0,
        }
    }
}

impl<'a> Iterator for QoiIterator<'a> {
    type Item = Rgb565;

    fn next(&mut self) -> Option<Self::Item> {
        if self.parsed_colors >= self.expected_colors {
            return None;
        }
        self.parsed_colors += 1;

        if self.repeat > 0 {
            self.repeat -= 1;
            return Some(self.last_pixel);
        }
        if let Some(byte) = self.inner.next() {
            let color = if byte == 0xff {
                let red = self.inner.next().unwrap();
                let green = self.inner.next().unwrap();
                let blue = self.inner.next().unwrap();
                let _alpha = self.inner.next().unwrap();
                Rgb565(red, green, blue)
            } else if byte == 0b11111110 {
                let red = self.inner.next()?;
                let green = self.inner.next()?;
                let blue = self.inner.next()?;
                Rgb565(red, green, blue)
            } else {
                let tag = (0b11000000 & byte) >> 6;
                let data = 0b00111111 & byte;

                if tag == 0 {
                    self.prev_pixels[data as usize]
                } else if tag == 1 {
                    let dr = ((0b110000 & data) >> 4) as i8 - 2;
                    let dg = ((0b001100 & data) >> 2) as i8 - 2;
                    let db = (0b000011 & data) as i8 - 2;
                    Rgb565(
                        (self.last_pixel.0 as i8).wrapping_add(dr) as u8,
                        (self.last_pixel.1 as i8).wrapping_add(dg) as u8,
                        (self.last_pixel.2 as i8).wrapping_add(db) as u8,
                    )
                } else if tag == 2 {
                    let second = self.inner.next()?;
                    let dg = data as i8 - 32;
                    let dr_dg = ((0b11110000 & second) >> 4) as i8 - 8;
                    let db_dg = (0b00001111 & second) as i8 - 8;
                    let dr = dr_dg + dg;
                    let db = db_dg + dg;

                    Rgb565(
                        (self.last_pixel.0 as i8 + dr) as u8,
                        (self.last_pixel.1 as i8 + dg) as u8,
                        (self.last_pixel.2 as i8 + db) as u8,
                    )
                } else if tag == 3 {
                    // QOI_OP_RUN
                    self.repeat = data;
                    self.last_pixel
                } else {
                    Rgb565::green()
                }
            };
            self.last_pixel = color;
            let hash =
                ((color.0 as u32 * 3 + color.1 as u32 * 5 + color.2 as u32 * 7 + 255 as u32 * 11)
                    % 64) as usize;

            self.prev_pixels[hash] = color;
            Some(color)
        } else {
            None
        }
    }
}