use alloc::sync::Arc;
use embassy_sync::{blocking_mutex::raw::CriticalSectionRawMutex, mutex::Mutex, pubsub::DynSubscriber};
use embassy_time::{Duration, Timer};
use embedded_graphics::pixelcolor::BinaryColor;
use figments::{liber8tion::interpolate::Fract8, mappings::embedded_graphics::Matrix2DSpace, prelude::{Coordinates, Rectangle}, render::Shader, surface::{Surface, SurfaceBuilder, Surfaces}};
use figments_render::output::Brightness;
use log::*;

use crate::{animation::Animation, events::{Personality, Prediction}, graphics::{display::DisplayControls, oled_ui::{OledUniforms, Screen}}};

#[cfg(feature="oled")]
pub type OledUiSurfacePool = figments::surfaces::buffered::BufferedSurfacePool<OledUniforms, Matrix2DSpace, BinaryColor>;

#[cfg(not(feature="oled"))]
pub type OledUiSurfacePool = figments::surfaces::null::NullBufferPool<OledUniforms, Matrix2DSpace, BinaryColor>;

type OledSurface = <OledUiSurfacePool as Surfaces>::Surface;

pub type LockedUniforms = Arc<Mutex<CriticalSectionRawMutex, OledUniforms>>;

pub struct OledUI<S: Surface + core::fmt::Debug> {
    overlay: S,
    controls: DisplayControls,
    uniforms: LockedUniforms
}

struct OverlayShader {}
impl Shader<OledUniforms, Matrix2DSpace, BinaryColor> for OverlayShader {
    fn draw(&self, _surface_coords: &Coordinates<Matrix2DSpace>, _uniforms: &OledUniforms) -> BinaryColor {
        BinaryColor::On
    }
}

impl<S: core::fmt::Debug + Surface<CoordinateSpace = Matrix2DSpace, Pixel = BinaryColor, Uniforms = OledUniforms>> OledUI<S> {
    pub fn new<SS: Surfaces<Surface = S>>(surfaces: &mut SS, controls: DisplayControls, uniforms: LockedUniforms) -> Self where SS::Error: core::fmt::Debug {
        Self {
            overlay: SurfaceBuilder::build(surfaces)
                .rect(Rectangle::everything())
                .shader(OverlayShader{})
                .opacity(Fract8::MIN)
                .finish().unwrap(),
            controls,
            uniforms
        }
    }

    pub async fn screen_transition(&mut self, next_screen: Screen) {
        const FADE_IN: Animation<Fract8> = Animation::new().from(Fract8::MIN).to(Fract8::MAX).duration(Duration::from_millis(300));
        const FADE_OUT: Animation<Fract8> = Animation::new().from(Fract8::MAX).to(Fract8::MIN).duration(Duration::from_millis(300));
        info!("Fading in to screen {next_screen:?}");
        FADE_IN.apply([&mut self.overlay]).await;
        {
            let mut locked = self.uniforms.lock().await;
            locked.current_screen = next_screen;
        }
        FADE_OUT.apply([&mut self.overlay]).await;
    }

    pub async fn on_event(&mut self, event: Prediction) {
        match event {
            Prediction::SetPersonality(personality) => {
                match personality {
                    Personality::Waking => {
                        warn!("Waking up OLED display");
                        self.controls.set_on(true);
                        self.screen_transition(Screen::Waking).await;
                        Timer::after_secs(1).await;
                        self.screen_transition(Screen::Home).await;
                    },
                    Personality::Sleeping => {
                        warn!("Putting OLED display to sleep");
                        self.screen_transition(Screen::Sleeping).await;
                        Timer::after_secs(1).await;
                        self.screen_transition(Screen::Blank).await;
                        self.controls.set_on(false);
                    },
                    _ => ()
                }
                self.with_uniforms(|state| { state.ui.personality = personality }).await;
            },
            Prediction::Velocity(v) => self.with_uniforms(|state| {state.ui.velocity = v;}).await,
            Prediction::Location(loc) => self.with_uniforms(|state| {state.ui.location = loc}).await,
            Prediction::Motion{ prev: _, next: motion } => self.with_uniforms(|state| {state.ui.motion = motion}).await,
            Prediction::SensorStatus(src, sensor_state, ) => self.with_uniforms(|state| {state.ui.sensor_states[src] = sensor_state}).await,
        }
    }

    pub async fn with_uniforms(&self, f: impl Fn(&mut OledUniforms)) {
        let mut locked = self.uniforms.lock().await;
        let mutref: &mut OledUniforms = &mut locked;
        f(mutref);
    }
}

#[embassy_executor::task]
pub async fn oled_ui(mut events: DynSubscriber<'static, Prediction>, mut ui: OledUI<OledSurface>) {
    
    ui.screen_transition(Screen::Bootsplash).await;
    Timer::after_secs(3).await;
    ui.screen_transition(Screen::Home).await;

    loop {
        ui.on_event(events.next_message_pure().await).await;
    }

}