#![cfg(feature="oled")]
use core::{cmp::min, ops::{BitAnd, Shr}};

use display_interface::DisplayError;
use embedded_graphics::prelude::*;
use esp_hal::{gpio::Output, i2c::master::I2c, Async};
use figments::{liber8tion::{interpolate::Fract8}, mappings::embedded_graphics::Matrix2DSpace, prelude::*};
use embedded_graphics::pixelcolor::BinaryColor;
use figments::pixels::AdditivePixelSink;
use figments_render::output::OutputAsync;
use ssd1306::{prelude::DisplayRotation, size::DisplaySize128x64, I2CDisplayInterface, Ssd1306Async};
use embassy_time::Delay;

use crate::graphics::display::DisplayControls;

#[derive(Clone, Copy, Debug)]
pub struct SsdPixel {
    byte: *mut u8,
    bit: u8,
    coords: Coordinates<Matrix2DSpace>
}

impl SsdPixel {
    fn set_pixel(&self, value: BinaryColor) {
        if value.is_on() {
            unsafe { *self.byte |= 1 << self.bit };
        } else {
            unsafe { *self.byte &= !(1 << self.bit) };
        }
    }

    fn get_pixel(&self) -> BinaryColor {
        if unsafe { *self.byte >> self.bit & 0x01 == 0x01 } {
            BinaryColor::On
        } else {
            BinaryColor::Off
        }
    }
}

const DITHER_MAP: [u16;15] = [
    0b1000_0000_0000_0000,
    0b1000_0000_0010_0000,
    0b1010_0000_0010_0000,
    0b1010_0000_1010_0000,
    0b1010_0100_1010_0000,
    0b1010_0100_1010_0001,
    0b1010_0101_1010_0001,
    0b1010_0101_1010_0101,
    0b1110_0101_1010_0101,
    0b1110_0101_1011_0101,
    0b1111_0101_1011_0101,
    0b1111_0101_1111_0101,
    0b1111_1101_1111_0101,
    0b1111_1101_1111_0111,
    0b1111_1111_1111_0111,
];

impl AdditivePixelSink<BinaryColor> for SsdPixel {
    fn add(&mut self, pixel: BinaryColor, opacity: Fract8) {
        match opacity {
            Fract8::MIN => (),
            Fract8::MAX => self.set_pixel(pixel),
            _ => {
                let dither_value = DITHER_MAP[opacity.to_raw() as usize / 17];
                let dither_x = self.coords.x % 4;
                let dither_y = self.coords.y % 4;
                if dither_value.shr(dither_x).shr(dither_y * 4).bitand(0x01) == 1 {
                    self.set_pixel(pixel);
                }
            }
        }
    }
}

pub struct SsdSampler<'a> {
    buf: &'a mut SsdOutput,
    rect: figments::prelude::Rectangle<Matrix2DSpace>,
    pos: figments::prelude::Coordinates<Matrix2DSpace>,
    pixel: Option<SsdPixel>
}

impl<'a> Iterator for SsdSampler<'a> {
    type Item = (Coordinates<Matrix2DSpace>, &'a mut SsdPixel);

    fn next(&mut self) -> Option<Self::Item> {
        if self.pos.x > min(127, self.rect.right()) {
            self.pos.x = self.rect.left();
            self.pos.y += 1;
        }

        if self.pos.y > min(63, self.rect.bottom()) {
            return None;
        }

        self.pixel = self.buf.get_pixel(self.pos);
        let old_pos = self.pos;
        self.pos.x += 1;

        let pixelref = unsafe { &mut *(self.pixel.as_mut().unwrap() as *mut SsdPixel) };

        Some((old_pos, pixelref))
    }
}

pub struct SsdOutput {
    pixbuf: [u8; 128 * 64 / 8], 
    target: ssd1306::Ssd1306Async<ssd1306::prelude::I2CInterface<esp_hal::i2c::master::I2c<'static, esp_hal::Async>>, ssd1306::prelude::DisplaySize128x64, ssd1306::mode::BasicMode>,
    controls: DisplayControls,
    is_on: bool,
    last_brightness: Fract8,
    reset_pin: Output<'static>
}

impl SsdOutput {
    pub fn get_pixel(&mut self, coords: Coordinates<Matrix2DSpace>) -> Option<SsdPixel> {
        if coords.x < 0 || coords.y < 0 || coords.x >= 128 || coords.y >= 64 {
            return None;
        }
        let idx = (coords.y / 8 * 128 + coords.x) as usize;
        let bit = (coords.y % 8) as u8;
        //info!("sample {coords:?} {idx}");
        let pixref = &mut self.pixbuf[idx] as *mut u8;
        Some(SsdPixel { byte: pixref, bit, coords })
    }

    pub async fn new(i2c: I2c<'static, Async>, reset_pin: Output<'static>, controls: DisplayControls) -> Self {
        let interface = I2CDisplayInterface::new(i2c);
        let target = Ssd1306Async::new(interface, DisplaySize128x64, DisplayRotation::Rotate0);

        Self {
            pixbuf: [0; 128 * 64 / 8],
            target,
            controls,
            last_brightness: Fract8::MAX,
            is_on: true,
            reset_pin
        }
    }

    pub async fn init(&mut self) -> Result<(), display_interface::DisplayError> {
        self.target.reset(&mut self.reset_pin, &mut Delay).await.map_err(|_| { DisplayError::BusWriteError })?;
        self.target.init_with_addr_mode(ssd1306::command::AddrMode::Horizontal).await
    }
}

impl<'a> Sample<'a, Matrix2DSpace> for SsdOutput {
    type Output = SsdPixel;
    
    fn sample(&mut self, rect: &figments::prelude::Rectangle<Matrix2DSpace>) -> impl Iterator<Item = (Coordinates<Matrix2DSpace>, &'a mut Self::Output)> {
        let bufref = unsafe {
            &mut *(self as *mut Self)
        };
        let rect = figments::geometry::Rectangle::new(
            figments::geometry::Coordinates::new(rect.left().clamp(0, 128), rect.top().clamp(0, 64)),
            figments::geometry::Coordinates::new(rect.right().clamp(0, 128), rect.bottom().clamp(0, 64)),
        );
        SsdSampler {
            buf: bufref,
            pos: rect.top_left,
            rect,
            pixel: None
        }
    }
}

impl DrawTarget for SsdOutput {
    type Color = BinaryColor;

    type Error = ();

    fn draw_iter<I>(&mut self, pixels: I) -> Result<(), Self::Error>
    where
        I: IntoIterator<Item = Pixel<Self::Color>> {
        for epix in pixels {
            self.get_pixel(epix.0.into()).unwrap().set_pixel(epix.1);
        }
        Ok(())
    }

    fn fill_contiguous<I>(&mut self, area: &embedded_graphics::primitives::Rectangle, colors: I) -> Result<(), Self::Error>
        where
            I: IntoIterator<Item = Self::Color>, {
        let rect = Rectangle::new(
            area.top_left.into(),
            area.bottom_right().unwrap().into()
        );
        for (color, (_coords, pix)) in colors.into_iter().zip(self.sample(&rect)) {
            pix.set_pixel(color);
        }
        Ok(())
    }

    fn fill_solid(&mut self, area: &embedded_graphics::primitives::Rectangle, color: Self::Color) -> Result<(), Self::Error> {
        let rect = Rectangle::new(
            area.top_left.into(),
            area.bottom_right().unwrap().into()
        );
        for (_coords, pix) in self.sample(&rect) {
            pix.set_pixel(color);
        }
        Ok(())
    }

    fn clear(&mut self, color: Self::Color) -> Result<(), Self::Error> {
        if color.is_on() {
            self.pixbuf.fill_with(|| { 255 });
        } else {
            self.pixbuf.fill_with(|| { 0 });
        }
        Ok(())
    }
}

impl OriginDimensions for SsdOutput {    
    fn size(&self) -> Size {
        Size::new(128, 64)
    }
}

impl<'a> OutputAsync<'a, Matrix2DSpace> for SsdOutput {
    type Error = DisplayError;

    type Controls = DisplayControls;

    async fn commit_async(&mut self)  -> Result<(), Self::Error> {
        let new_brightness = self.controls.brightness();
        let new_power = self.controls.is_on();
        if self.is_on != new_power {
            self.target.set_display_on(self.controls.is_on()).await.unwrap();
            self.is_on = new_power;
        }
        if self.last_brightness != new_brightness {
            self.target.set_brightness(ssd1306::prelude::Brightness::custom(1, new_brightness.to_raw())).await.unwrap();
            self.last_brightness = new_brightness;
        }
        
        self.target.draw(&self.pixbuf).await
    }

    fn controls(&mut self) -> Option<&mut Self::Controls> {
        Some(&mut self.controls)
    }
}