use embassy_sync::{blocking_mutex::raw::CriticalSectionRawMutex, signal::Signal, watch::{Receiver, Watch}};
use figments::{liber8tion::interpolate::Fract8, prelude::*};
use core::{fmt::Debug, ops::Mul, sync::atomic::{AtomicBool, AtomicU8}};
use alloc::sync::Arc;

use figments_render::{
    gamma::{GammaCurve, WithGamma}, output::{Brightness, GammaCorrected, Output, OutputAsync}, power::AsMilliwatts, smart_leds::PowerManagedWriter
};
use smart_leds::{SmartLedsWrite, SmartLedsWriteAsync};

pub const NUM_PIXELS: usize = 178;
#[derive(Debug)]
pub struct BikeOutput<T, Color> {
    pixbuf: [Color; NUM_PIXELS],
    writer: PowerManagedWriter<T>,
    controls: DisplayControls
}

impl<T, Color> GammaCorrected for BikeOutput<T, Color> {
    fn set_gamma(&mut self, gamma: GammaCurve) {
        self.writer.controls().set_gamma(gamma);
    }
}

impl<T, Color: Default + Copy> BikeOutput<T, Color> {
    pub fn new(target: T, max_mw: u32, controls: DisplayControls) -> Self {
        Self {
            pixbuf: [Default::default(); NUM_PIXELS],
            writer: PowerManagedWriter::new(target, max_mw),
            controls
        }
    }

    pub fn blank(&mut self) {
        self.pixbuf = [Default::default(); NUM_PIXELS];
    }
}

impl<'a, T: SmartLedsWrite + 'a> Output<'a, SegmentSpace> for BikeOutput<T, T::Color> where T::Color: AsMilliwatts + Copy + Mul<Fract8, Output = T::Color> + WithGamma + Default + Debug + 'a + 'static {
    type Error = T::Error;

    type Controls = DisplayControls;

    fn commit(&mut self)  -> Result<(), Self::Error> {
        self.writer.controls().set_brightness(self.controls.brightness());
        self.writer.controls().set_on(self.controls.is_on());
        critical_section::with(|_| {
            self.writer.write(&self.pixbuf)
        })
    }

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

impl<'a, T: SmartLedsWriteAsync + 'a> OutputAsync<'a, SegmentSpace> for BikeOutput<T, T::Color> where T::Color: 'static + Copy + Mul<Fract8, Output = T::Color> + WithGamma + Default + Debug + AsMilliwatts + 'a {
    async fn commit_async(&mut self) -> Result<(), T::Error> where T: SmartLedsWriteAsync {
        self.writer.controls().set_brightness(self.controls.brightness());
        self.writer.controls().set_on(self.controls.is_on());
        // TODO: We should grab the power used here and somehow feed it back into the telemetry layer, probably just via another atomic u32
        self.writer.write_async(&self.pixbuf).await
    }
    
    type Error = T::Error;
    type Controls = DisplayControls;
    
    fn controls(&mut self) -> Option<&mut Self::Controls> {
        Some(&mut self.controls)
    }
}

impl<'a, T, Color> Sample<'a, SegmentSpace> for BikeOutput<T, Color> where Color: 'a, [Color; NUM_PIXELS]: Sample<'a, SegmentSpace, Output = Color> {
    type Output = Color;

    fn sample(&mut self, rect: &Rectangle<SegmentSpace>) -> impl Iterator<Item = (Coordinates<SegmentSpace>, &'a mut Self::Output)> {
        self.pixbuf.sample(rect)
    }
}

const STRIP_MAP: [Segment; 6] = [
    Segment { start: 0, length: 28 },
    Segment { start: 28, length: 17 },
    Segment { start: 45, length: 14 },
    Segment { start: 59, length: 17 },
    Segment { start: 76, length: 14 },
    Segment { start: 90, length: 88 }
];

pub struct Segment {
    start: usize,
    length: usize,
}

impl<'a, T> Sample<'a, SegmentSpace> for [T; NUM_PIXELS] where T: 'static {
    type Output = T;

    fn sample(&mut self, rect: &Rectangle<SegmentSpace>) -> impl Iterator<Item = (Coordinates<SegmentSpace>, &'a mut Self::Output)> {
        let bufref = unsafe {
            &mut *(self as *mut [T; NUM_PIXELS])
        };
        SegmentIter {
            pixbuf: bufref,
            cur: rect.top_left,
            end: rect.bottom_right
        }
    }
}

#[derive(Clone, Copy, Default, Debug)]
pub struct SegmentSpace {}
impl CoordinateSpace for SegmentSpace {
    type Data = usize;
}

pub struct SegmentIter<'a, Format> {
    pixbuf: &'a mut [Format; NUM_PIXELS],
    cur: Coordinates<SegmentSpace>,
    end: Coordinates<SegmentSpace>,
}

impl<'a, Format: Debug> Debug for SegmentIter<'a, Format> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("BikeIter").field("cur", &self.cur).field("end", &self.end).finish()
    }
}

impl<'a, Format> Iterator for SegmentIter<'a, Format> {
    type Item = (Coordinates<SegmentSpace>, &'a mut Format);

    fn next(&mut self) -> Option<Self::Item> {
        if self.cur.y > self.end.y || self.cur.y >= STRIP_MAP.len() {
            return None;
        }

        let this_strip = &STRIP_MAP[self.cur.y];

        let offset = this_strip.start + self.cur.x;
        let pixel_coords = Coordinates::new(self.cur.x, self.cur.y);

        self.cur.x += 1;

        if self.cur.x >= this_strip.length || self.cur.x > self.end.x {
            self.cur.x = 0;
            self.cur.y += 1;
        }

        let bufref = unsafe {
            &mut *(self.pixbuf as *mut [Format; NUM_PIXELS])
        };

        Some((pixel_coords, &mut bufref[offset]))
    }
}

#[derive(Default, Debug)]
pub struct Uniforms {
    pub frame: usize,
    pub primary_color: Hsv
}

struct ControlData {
    on: AtomicBool,
    brightness: AtomicU8
}

impl Default for ControlData {
    fn default() -> Self {
        Self {
            on: AtomicBool::new(true),
            brightness: AtomicU8::new(255)
        }
    }
}

// A watch that indicates whether or not the rendering engine is running. If the display is off or sleeping, this will be false.
static RENDER_IS_RUNNING: Watch<CriticalSectionRawMutex, bool, 7> = Watch::new();

// TODO: Implement something similar for a system-wide sleep mechanism
pub struct DisplayControls {
    data: Arc<ControlData>,
    display_is_on: Arc<Signal<CriticalSectionRawMutex, bool>>,
    render_run_receiver: Receiver<'static, CriticalSectionRawMutex, bool, 7>
}

impl Clone for DisplayControls {
    fn clone(&self) -> Self {
        Self {
            data: Arc::clone(&self.data),
            display_is_on: Arc::clone(&self.display_is_on),
            render_run_receiver: RENDER_IS_RUNNING.receiver().expect("Could not create enough render running receivers")
        }
    }
}

impl DisplayControls {
    pub fn is_on(&self) -> bool {
        self.data.on.load(core::sync::atomic::Ordering::Relaxed)
    }

    pub fn brightness(&self) -> Fract8 {
        Fract8::from_raw(self.data.brightness.load(core::sync::atomic::Ordering::Relaxed))
    }

    pub async fn wait_until_display_is_on(&self) {
        while !self.display_is_on.wait().await { log::info!("wait for display") }
        log::trace!("display says on!");
    }

    pub fn notify_render_is_running(&mut self, value: bool) {
        log::trace!("render is running!");
        RENDER_IS_RUNNING.sender().send(value);
    }

    pub async fn wait_until_render_is_running(&mut self) {
        while !self.render_run_receiver.changed().await { log::info!("wait for render run") }
        log::trace!("render says run!");
    }
}

impl GammaCorrected for DisplayControls {
    fn set_gamma(&mut self, _gamma: GammaCurve) {
        todo!()
    }
}

impl Brightness for DisplayControls {
    fn set_brightness(&mut self, brightness: Fract8) {
        self.data.brightness.store(brightness.to_raw(), core::sync::atomic::Ordering::Relaxed);
    }

    fn set_on(&mut self, is_on: bool) {
        self.data.on.store(is_on, core::sync::atomic::Ordering::Relaxed);
        log::trace!("display is on {is_on}");
        self.display_is_on.signal(is_on);
    }
}

impl core::fmt::Debug for DisplayControls {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> Result<(), core::fmt::Error> {
        f.debug_struct("DisplayControls")
            .field("on", &self.data.on)
            .field("brightness", &self.data.brightness)
            .field("render_pause_signaled", &self.display_is_on.signaled())
            .finish()
    }
}

impl Default for DisplayControls {
    fn default() -> Self {
        Self {
            data: Default::default(),
            display_is_on: Default::default(),
            render_run_receiver: RENDER_IS_RUNNING.receiver().unwrap()
        }
    }
}