use core::{cell::{Cell, RefCell}, cmp::min, f32::consts::PI, fmt::Binary, ops::DerefMut};

use alloc::{format, sync::Arc};
use display_interface::DisplayError;
use embassy_executor::{raw, Spawner};
use embassy_sync::{blocking_mutex::raw::{CriticalSectionRawMutex, NoopRawMutex}, mutex::Mutex, pubsub::DynSubscriber};
use embassy_time::{Delay, Duration, Instant, Timer};
use embedded_graphics::{image::Image, mono_font::{ascii::{FONT_6X10, FONT_10X20, FONT_9X18_BOLD}, MonoTextStyleBuilder}, pixelcolor::BinaryColor, prelude::*, primitives::{Line, PrimitiveStyleBuilder, Rectangle, StyledDrawable}, text::{Alignment, Baseline, Text}};
use esp_hal::{i2c::master::I2c, Async, gpio::Output};
use figments::{liber8tion::{interpolate::{ease_in_out_quad, scale8, Fract8}, trig::{cos8, sin8}}, mappings::{embedded_graphics::{EmbeddedGraphicsSampler, Matrix2DSpace}, linear::LinearSpace}, pixels::{PixelBlend, PixelSink}, prelude::{CoordinateSpace, Coordinates, Fract8Ops}, render::{RenderSource, Sample}, surface::{BufferedSurface, BufferedSurfacePool, Surface, SurfaceBuilder, Surfaces}};
use figments_render::output::{Brightness, NullControls, OutputAsync};
use futures::{join, FutureExt};
use nalgebra::{Matrix, Vector2};
use ssd1306::{mode::{BufferedGraphicsModeAsync, DisplayConfigAsync}, prelude::{DisplayRotation, I2CInterface}, size::DisplaySize128x64, I2CDisplayInterface, Ssd1306Async};
use log::*;
use embedded_graphics::mono_font::MonoTextStyle;
use embedded_graphics::primitives::PrimitiveStyle;
use nalgebra::ComplexField;

use crate::{animation::{AnimatedSurface, Animation}, backoff::Backoff, ego::engine::MotionState, events::{Notification, Prediction, Scene, SensorSource, Telemetry}, graphics::ssd1306::{SsdControls, SsdOutput, SsdPixel}, tasks::ui};

mod images {
    use embedded_graphics::{
        image::ImageRaw,
        pixelcolor::BinaryColor
    };
    include!("../../target/images.rs");
}

#[derive(Default, Debug)]
struct OledUI {
    scene: Scene,
    motion: MotionState,
    brakelight: bool,
    headlight: bool,
    gps_online: bool,
    imu_online: bool,
    velocity: f32,
    location: Vector2<f64>,
    sleep: bool,
}

#[derive(Default, Debug)]
struct OledUniforms {
    ui: OledUI,
    frame: usize,
    current_screen: Screen,
    next_screen: Screen
}

#[derive(Debug, Default, Clone, Copy)]
enum Screen {
    #[default]
    Blank,
    Bootsplash,
    Home,
    Sleeping,
    Waking
}

impl Screen {
    pub fn draw_screen<'a, T>(&self, sampler: &mut EmbeddedGraphicsSampler<'a, T>, state: &OledUniforms) where T: Sample<'a, Matrix2DSpace>, T::Output: PixelSink<BinaryColor> {
        match self {
            Screen::Blank => (),
            Screen::Bootsplash => {
                Image::new(&images::BOOT_LOGO, Point::zero()).draw(sampler).unwrap();
                const SPARKLE_COUNT: i32 = 8;
                for n in 0..SPARKLE_COUNT {
                    let sparkle_center = Point::new(128u8.scale8(cos8(state.frame.wrapping_mul(n as usize))) as i32, 64u8.scale8(sin8(state.frame.wrapping_mul(n as usize) as u8)) as i32);
                    let offset = (state.frame / 2 % 32) as i32 - 16;
                    let rotation = PI * 2.0 * (sin8(state.frame) as f32 / 255.0);
                    let normal = Point::new((rotation.cos() * offset as f32) as i32, (rotation.sin() * offset as f32) as i32);
                    let cross_normal = Point::new((rotation.sin() * offset as f32) as i32, (rotation.cos() * offset as f32) as i32);
                    // Draw horizontal
                    Line::new(sparkle_center - normal, sparkle_center + normal)
                        .draw_styled(&SPARKLE_STYLE, sampler).unwrap();
                    // Draw vertical
                    Line::new(sparkle_center - cross_normal, sparkle_center + cross_normal)
                        .draw_styled(&SPARKLE_STYLE, sampler).unwrap();
                }
            },
            Screen::Home => {
                // Status bar
                // Sensor indicators
                let gps_img = if state.ui.gps_online {
                    &images::GPS_ON
                } else {
                    &images::GPS_OFF
                };
                let imu_img = if state.ui.imu_online {
                    &images::IMU_ON
                } else {
                    &images::IMU_OFF
                };
                Image::new(gps_img, Point::zero()).draw(sampler).unwrap();
                Image::new(imu_img, Point::new((gps_img.size().width + 2) as i32, 0)).draw(sampler).unwrap();

                #[cfg(feature="demo")]
                Text::with_alignment("Demo", Point::new(128, 10), TEXT_STYLE, Alignment::Right)
                    .draw(sampler)
                    .unwrap();

                #[cfg(feature="simulation")]
                Text::with_alignment("Sim", Point::new(128, 10), TEXT_STYLE, Alignment::Right)
                    .draw(sampler)
                    .unwrap();

                // Separates the status bar from the UI
                Line::new(Point::new(0, 18), Point::new(128, 18)).draw_styled(&INACTIVE_STYLE, sampler).unwrap();

                // Speed display at the top
                Text::with_alignment(&format!("{}", state.ui.velocity), Point::new(128 / 2, 12), SPEED_STYLE, Alignment::Center)
                    .draw(sampler)
                    .unwrap();

                // The main UI content
                Image::new(&images::BIKE, Point::new((128 / 2 - images::BIKE.size().width / 2) as i32, 24)).draw(sampler).unwrap();

                let headlight_img = if state.ui.headlight {
                    &images::HEADLIGHT_ON
                } else {
                    &images::HEADLIGHT_OFF
                };
                let brakelight_img = if state.ui.brakelight {
                    &images::BRAKELIGHT_ON
                } else {
                    &images::BRAKELIGHT_OFF
                };
                Image::new(headlight_img, Point::new(((128 / 2 - images::BIKE.size().width / 2) - 18) as i32, 28)).draw(sampler).unwrap();
                Image::new(brakelight_img, Point::new(((128 / 2 + images::BIKE.size().width / 2) + 2) as i32, 28)).draw(sampler).unwrap();

                // TODO: Replace the state texts with cute animations or smth
                // Current prediction from the motion engine
                Text::with_alignment(&format!("{:?}", state.ui.motion), Point::new(128 / 2, 64 - 3), TEXT_STYLE, Alignment::Center)
                    .draw(sampler)
                    .unwrap();
                
                // Current scene in the UI
                Text::with_alignment(&format!("{:?}", state.ui.scene), Point::new(128 / 2, 64 - 13), TEXT_STYLE, Alignment::Center)
                    .draw(sampler)
                    .unwrap();
            },
            Screen::Sleeping => {
                Text::with_alignment("so eepy Zzz...", Point::new(128 / 2, 64 / 2), LOGO_STYLE, Alignment::Center)
                    .draw(sampler)
                    .unwrap();
            },
            Screen::Waking => {
                Text::with_alignment("OwO hewwo again :D", Point::new(128 / 2, 64 / 2), LOGO_STYLE, Alignment::Center)
                    .draw(sampler)
                    .unwrap();
            }
        }
    }
}

impl RenderSource<OledUniforms, Matrix2DSpace, BinaryColor, SsdPixel> for Screen {
    fn render_to<'a, Smp>(&'a self, output: &'a mut Smp, uniforms: &OledUniforms)
        where 
            Smp: Sample<'a, Matrix2DSpace, Output = SsdPixel> {
        let mut sampler = EmbeddedGraphicsSampler(output, embedded_graphics::primitives::Rectangle::new(Point::new(0, 0), Size::new(128, 64)));
        self.draw_screen(&mut sampler, uniforms);
    }
}

const SPARKLE_STYLE: PrimitiveStyle<BinaryColor> = PrimitiveStyleBuilder::new()
    .stroke_color(BinaryColor::On)
    .stroke_width(2)
    .build();

const TEXT_STYLE: MonoTextStyle<BinaryColor> = MonoTextStyleBuilder::new()
    .font(&FONT_6X10)
    .text_color(BinaryColor::On)
    .build();

const LOGO_STYLE: MonoTextStyle<BinaryColor> = MonoTextStyleBuilder::new()
    .font(&FONT_9X18_BOLD)
    .text_color(BinaryColor::On)
    .build();

const SPEED_STYLE: MonoTextStyle<BinaryColor> = MonoTextStyleBuilder::new()
    .font(&FONT_10X20)
    .text_color(BinaryColor::On)
    .build();

const DRAW_STYLE: PrimitiveStyle<BinaryColor> = PrimitiveStyleBuilder::new()
    .fill_color(BinaryColor::On)
    .build();

const INACTIVE_STYLE: PrimitiveStyle<BinaryColor> = PrimitiveStyleBuilder::new()
    .fill_color(BinaryColor::Off)
    .stroke_color(BinaryColor::On)
    .stroke_width(2)
    .build();

#[embassy_executor::task]
async fn oled_render(mut output: SsdOutput, surfaces: BufferedSurfacePool<OledUniforms, Matrix2DSpace, BinaryColor>, uniforms: Arc<Mutex<NoopRawMutex, OledUniforms>>) {
    Backoff::from_secs(1).forever().attempt::<_, (), DisplayError>(async || {
        const FPS: u64 = 30;
        const RENDER_BUDGET: Duration = Duration::from_millis(1000 / FPS);

        const ANIMATION_TPS: u64 = 30;
        const ANIMATION_FRAME_TIME: Duration = Duration::from_millis(1000 / ANIMATION_TPS);
        info!("Starting Oled renderer");
        loop {
            let start = Instant::now();
            {
                let mut locked = uniforms.lock().await;
                output.clear(BinaryColor::Off).unwrap();
                locked.frame = (Instant::now().as_millis() / ANIMATION_FRAME_TIME.as_millis()) as usize;
                locked.current_screen.render_to(&mut output, &locked);
                surfaces.render_to(&mut output, &locked);
            }
            output.commit_async().await?;
            let frame_time = Instant::now() - start;
            if frame_time < RENDER_BUDGET {
                Timer::after(RENDER_BUDGET - frame_time).await;
            } else {
                //warn!("OLED Frame took too long to render! {}ms", frame_time.as_millis());
            }
        }
    }).await.unwrap();
}

async fn screen_transition(overlay: &mut BufferedSurface<OledUniforms, Matrix2DSpace, BinaryColor>, uniforms: &mut Arc<Mutex<NoopRawMutex, OledUniforms>>, next_screen: Screen) {
    const FADE_IN: Animation = Animation::new().from(0).to(255).duration(Duration::from_millis(300));
    const FADE_OUT: Animation = Animation::new().from(255).to(0).duration(Duration::from_millis(300));
    info!("Fading in to screen {next_screen:?}");
    FADE_IN.apply(overlay).await;
    {
        let mut locked = uniforms.lock().await;
        locked.current_screen = next_screen;
    }
    FADE_OUT.apply(overlay).await;
}

#[embassy_executor::task]
async fn oled_ui(mut events: DynSubscriber<'static, Telemetry>, mut overlay: BufferedSurface<OledUniforms, Matrix2DSpace, BinaryColor>, mut uniforms: Arc<Mutex<NoopRawMutex, OledUniforms>>, mut controls: SsdControls) {
    
    screen_transition(&mut overlay, &mut uniforms, Screen::Bootsplash).await;
    Timer::after_secs(3).await;
    screen_transition(&mut overlay, &mut uniforms, Screen::Home).await;

    loop {
        let evt = events.next_message_pure().await;
        //info!("Oled UI event: {evt:?}");
        match evt {
            Telemetry::Notification(Notification::Sleep) => {
                screen_transition(&mut overlay, &mut uniforms, Screen::Sleeping).await;
                Timer::after_secs(1).await;
                screen_transition(&mut overlay, &mut uniforms, Screen::Blank).await;
                controls.set_on(false);
                //ui_state.sleep = true
            },
            Telemetry::Notification(Notification::WakeUp) => {
                controls.set_on(true);
                screen_transition(&mut overlay, &mut uniforms, Screen::Waking).await;
                Timer::after_secs(1).await;
                screen_transition(&mut overlay, &mut uniforms, Screen::Home).await;
                //ui_state.sleep = false
            },
            _ => ()
        }
        let mut locked = uniforms.lock().await;
        let ui_state = &mut locked.ui;
        match evt {
            Telemetry::Prediction(Prediction::Velocity(v)) => ui_state.velocity = v,
            Telemetry::Prediction(Prediction::Location(loc)) => ui_state.location = loc,
            Telemetry::Prediction(Prediction::Motion(motion)) => ui_state.motion = motion,
            Telemetry::Notification(Notification::SceneChange(scene)) => ui_state.scene = scene,
            Telemetry::Notification(Notification::SetBrakelight(b)) => ui_state.brakelight = b,
            Telemetry::Notification(Notification::SetHeadlight(b)) => ui_state.headlight = b,
            Telemetry::Notification(Notification::SensorOffline(SensorSource::IMU)) => ui_state.imu_online = false,
            Telemetry::Notification(Notification::SensorOnline(SensorSource::IMU)) => ui_state.imu_online = true,
            Telemetry::Notification(Notification::SensorOffline(SensorSource::GPS)) => ui_state.gps_online = false,
            Telemetry::Notification(Notification::SensorOnline(SensorSource::GPS)) => ui_state.gps_online = true,
            _ => ()
        }
    }
    /*for (coords, pix) in output.sample(&figments::prelude::Rectangle::everything()) {
        pix.set(&BinaryColor::Off);
    }
    Text::with_alignment("Foo", Point::new(128 / 2, 64 / 2), LOGO_STYLE, Alignment::Center).draw(&mut output).unwrap();
    output.commit_async().await.unwrap();
    Timer::after_secs(1).await;

    let painter = UiPainter { current_screen: Screen::Bootsplash, next_screen: Screen::Bootsplash };
    painter.render_to(&mut output, &OledUniforms { frame: 0, ui: ui_state });*/

}

#[embassy_executor::task]
pub async fn oled_task(i2c: I2c<'static, Async>,  mut reset_pin: Output<'static>, events: DynSubscriber<'static, Telemetry>) {
    let interface = I2CDisplayInterface::new(i2c);
    let mut display = Ssd1306Async::new(interface, DisplaySize128x64, DisplayRotation::Rotate0);
    display.reset(&mut reset_pin, &mut Delay).await.unwrap();
    display.init_with_addr_mode(ssd1306::command::AddrMode::Horizontal).await.unwrap();

    let output = SsdOutput::new(display);
    let mut surfaces = BufferedSurfacePool::default();
    let uniforms = Default::default();

    let sfc = SurfaceBuilder::build(&mut surfaces).shader(|coords: &Coordinates<Matrix2DSpace>, uniforms: &OledUniforms| {
        BinaryColor::On
    }).opacity(0).finish().unwrap();

    info!("Starting OLED display tasks!");
    let spawner = Spawner::for_current_executor().await;
    let ui_controls = output.controls().unwrap().clone();
    spawner.must_spawn(oled_render(output, surfaces, Arc::clone(&uniforms)));
    spawner.must_spawn(oled_ui(events, sfc, uniforms, ui_controls));
}