#![no_std]
#![no_main]
#![deny(
    clippy::mem_forget,
    reason = "mem::forget is generally not safe to do with esp_hal types, especially those \
    holding buffers for the duration of a data transfer."
)]

use core::cell::RefCell;
use core::ptr::read;

use alloc::rc::Rc;
use alloc::string::String;
use embassy_embedded_hal::shared_bus::asynch::i2c::I2cDevice;
use embassy_executor::{SendSpawner, Spawner};
use embassy_sync::blocking_mutex::raw::{CriticalSectionRawMutex, NoopRawMutex};
use embassy_sync::mutex::Mutex;
use embassy_sync::pubsub::{DynPublisher, DynSubscriber, PubSubChannel, Subscriber};
use embassy_time::{Duration, Instant, Timer, WithTimeout};
use embedded_hal_async::i2c::{I2c as _};
use esp_backtrace as _;
use embassy_time::Delay;
use esp_hal::i2c::master::{Config, I2c};
use esp_hal::interrupt::software::SoftwareInterruptControl;
use esp_hal::{Async, Blocking};
use esp_hal::clock::CpuClock;
use esp_hal::rmt::{ChannelCreator, Rmt};
use esp_hal::time::Rate;
use esp_hal::timer::systimer::SystemTimer;
use esp_hal::timer::timg::TimerGroup;
use esp_hal_embassy::InterruptExecutor;
use esp_hal_smartled::{buffer_size_async, smart_led_buffer, SmartLedsAdapter, SmartLedsAdapterAsync};
use esp_hal::peripherals::GPIO5;
use figments::liber8tion::trig::sin8;
use figments::prelude::*;
use figments::smart_leds::BrightnessWriter;
use log::{info, warn};
use mpu6050_dmp::address::Address;
use mpu6050_dmp::calibration::CalibrationParameters;
use mpu6050_dmp::error_async::Error;
use mpu6050_dmp::sensor_async::Mpu6050;
use nmea::Nmea;
use static_cell::StaticCell;

extern crate alloc;

// This creates a default app-descriptor required by the esp-idf bootloader.
// For more information see: <https://docs.espressif.com/projects/esp-idf/en/stable/esp32/api-reference/system/app_image_format.html#application-description>
esp_bootloader_esp_idf::esp_app_desc!();

static I2C_BUS: StaticCell<Mutex<CriticalSectionRawMutex, I2c<'static, Async>>> = StaticCell::new();
static EVENT_BUS: StaticCell<PubSubChannel<NoopRawMutex, Notification, 4, 4, 4>> = StaticCell::new();

#[esp_hal_embassy::main]
async fn main(spawner: Spawner) {
    // generator version: 0.5.0

    esp_println::logger::init_logger_from_env();

    let config = esp_hal::Config::default().with_cpu_clock(CpuClock::max());
    let peripherals = esp_hal::init(config);

    esp_alloc::heap_allocator!(size: 64 * 1024);

    let timer0 = SystemTimer::new(peripherals.SYSTIMER);
    esp_hal_embassy::init(timer0.alarm0);

    info!("Embassy initialized!");

    let events = EVENT_BUS.init(PubSubChannel::<NoopRawMutex, Notification, 4, 4, 4>::new());

    let rng = esp_hal::rng::Rng::new(peripherals.RNG);
    let timer1 = TimerGroup::new(peripherals.TIMG0);
    let wifi_init =
        esp_wifi::init(timer1.timer0, rng).expect("Failed to initialize WIFI/BLE controller");
    let (mut _wifi_controller, _interfaces) = esp_wifi::wifi::new(&wifi_init, peripherals.WIFI)
        .expect("Failed to initialize WIFI controller");

    let frequency: Rate = Rate::from_mhz(80);
    let rmt = Rmt::new(peripherals.RMT, frequency)
        .expect("Failed to initialize RMT").into_async();
    let rmt_channel = rmt.channel0;

    info!("Launching render task");
    //spawner.must_spawn(render(events.subscriber().unwrap(), rmt_channel, peripherals.GPIO5));

    info!("Launching i2c task");
    let i2c = I2c::new(peripherals.I2C1, Config::default().with_frequency(Rate::from_khz(400))).unwrap().with_scl(peripherals.GPIO36).with_sda(peripherals.GPIO33).into_async();
    spawner.must_spawn(i2c_task(events.dyn_publisher().unwrap(), spawner, i2c));
}

async fn mpu_init(sensor: &mut Mpu6050<I2cDevice<'static, CriticalSectionRawMutex, I2c<'static, Async>>>) -> Result<(), Error<I2cDevice<'static, CriticalSectionRawMutex, I2c<'static, Async>>>> {
    let mut delay = Delay;
    let backoff = Backoff::from_millis(3);
    info!("Resetting MPU");
    backoff.attempt(async || {sensor.reset(&mut delay).await}).await?;
    info!("Uploading DMP firmware");
    backoff.attempt(async || {sensor.initialize_dmp(&mut delay).await}).await?;
    info!("Calibrating");
    backoff.attempt(async || {sensor.calibrate(&mut delay, &CalibrationParameters::new(
        mpu6050_dmp::accel::AccelFullScale::G2,
        mpu6050_dmp::gyro::GyroFullScale::Deg2000,
        mpu6050_dmp::calibration::ReferenceGravity::ZN
    )).await}).await?;
    info!("Configuring sample rate");
    backoff.attempt(async || {sensor.set_sample_rate_divider(255).await}).await
}

#[embassy_executor::task]
async fn i2c_task(events: DynPublisher<'static, Notification>, spawner: Spawner, i2c: I2c<'static, Async>) {
    let i2c_bus = I2C_BUS.init(Mutex::new(i2c));

    spawner.spawn(mpu_task(I2cDevice::new(i2c_bus))).unwrap();
    spawner.spawn(gps_task(events, I2cDevice::new(i2c_bus))).unwrap();
}

#[embassy_executor::task]
async fn mpu_task(bus: I2cDevice<'static, CriticalSectionRawMutex, I2c<'static, Async>>) {
    let backoff = Backoff::from_millis(5);
    let busref = RefCell::new(Some(bus));

    backoff.forever().attempt::<_, (), ()>(async || {
        let mut sensor = backoff.attempt(async || {
            info!("Initializing MPU");
            match Mpu6050::new(busref.replace(None).unwrap(), Address::default()).await.map_err(|e| { e.i2c }) {
                Err(i2c) => {
                    busref.replace(Some(i2c));
                    Err(())
                },
                Ok(mut sensor) => {
                    match backoff.attempt(async || { mpu_init(&mut sensor).await }).await {
                        Err(_) => {
                            busref.replace(Some(sensor.release()));
                            Err(())
                        },
                        Ok(_) => Ok(sensor)
                    }
                }
            }
        }).await?;
        info!("MPU is ready!");
        let sensor_ref = &mut sensor;
        loop {
            match backoff.attempt(async || { sensor_ref.motion6().await }).await {
                Ok((accel_data, gyro_data)) => {
                    info!("Accel x={} y={} z={}", accel_data.x() as i32, accel_data.y() as i32, accel_data.z() as i32);
                    info!("Gyro x={} y={} z={}", gyro_data.x() as i32, gyro_data.y() as i32, gyro_data.z() as i32);
                    Timer::after_millis(5000).await;
                },
                Err(e) => {
                    warn!("Failed to read MPU motion data! {e:?}");
                    busref.replace(Some(sensor.release()));
                    return Err(());
                }
            };
        }
    }).await.unwrap();
}

const GPS_TEST_DATA: &str = include_str!("../test.nmea");

#[embassy_executor::task]
async fn gps_task(events: DynPublisher<'static, Notification>, mut i2c_bus: I2cDevice<'static, CriticalSectionRawMutex, I2c<'static, Async>>) {
    Backoff::from_secs(5).forever().attempt(async || {
        info!("Initializing GPS");
        // Enable a bunch of data? idk
        let bytes = "$PMTK314,1,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0*28\r\n";
        i2c_bus.write(0x10, bytes.as_bytes()).await?;

        // 1hz updates
        let bytes = "$PMTK220,1000*1F\r\n";
        i2c_bus.write(0x10, bytes.as_bytes()).await?;

        // 1hz position fix
        let bytes = "$PMTK300,1000,0,0,0,0*1C\r\n";
        i2c_bus.write(0x10, bytes.as_bytes()).await?;

        // Antenna updates
        let bytes = "$PGCMD,33,1*6C\r\n";
        i2c_bus.write(0x10, bytes.as_bytes()).await
    }).await.unwrap();

    let mut strbuf = String::new();

    let mut parser = Nmea::default();
    let mut parsing = false;
    let mut has_lock = false;
    //let mut iter = GPS_TEST_DATA.as_bytes().iter();
    info!("GPS is ready!");
    loop {
        let mut buf = [0; 1];
        i2c_bus.read(0x10, &mut buf).await.unwrap();
        //buf[0] = *(iter.next().unwrap());
        if (buf[0] as char == '\n' || buf[0] as char == '\r') && !strbuf.is_empty() {
            if let Ok(result) = parser.parse(&strbuf) {
                if parser.fix_type.is_some() != has_lock {
                    has_lock = parser.fix_type.is_some();
                    if has_lock {
                        events.publish(Notification::GPSAcquired).await;
                    } else {
                        events.publish(Notification::GPSLost).await;
                    }
                }
                log::info!("nmea={result:?} raw={strbuf:?}");
                log::debug!("nmea={parser:?}");
                log::info!("speed={:?} altitude={:?} lat={:?} lng={:?} fix={:?}", parser.speed_over_ground, parser.altitude, parser.latitude, parser.longitude, parser.fix_type);
                for sat in parser.satellites() {
                    info!("\t{} snr={:?} prn={:?}", sat.gnss_type(), sat.snr(), sat.prn())
                }
            } else {
                log::warn!("Unhandled NMEA {strbuf:?}");
            }
            strbuf = String::new();
            parsing = false;
            // Update frequency is 1hz, so we should never get an update faster than once per second
            Timer::after_secs(1).await;
        } else if strbuf.is_empty() && (buf[0] as char == '$' || buf[0] as char == '!') {
            parsing = true;
            strbuf.push(buf[0] as char);
            Timer::after_millis(3).await;
        } else if parsing {
            strbuf.push(buf[0] as char);
            Timer::after_millis(3).await;
        } else {
            // If there is no data ready for some reason, wait 500ms, which should place us at least somewhere after the next data frame is ready to read.
            Timer::after_millis(500).await;
        }
    }
}

#[derive(Clone, Copy, Debug)]
struct Backoff {
    delay: Duration,
    attempts: Attempts,
}

#[derive(Clone, Copy, Debug)]
enum Attempts {
    Finite(u64),
    Forever
}

impl Iterator for Attempts {
    type Item = u64;

    fn next(&mut self) -> Option<Self::Item> {
        match *self {
            Attempts::Forever => Some(0),
            Attempts::Finite(0) => None,
            Attempts::Finite(n) => {
                *self = Attempts::Finite(n - 1);
                Some(n)
            }
        }
    }
}

impl Backoff {
    pub fn from_millis(millis: u64) -> Self {
        Self {
            delay: Duration::from_millis(millis),
            attempts: Attempts::Finite(3),
        }
    }

    pub fn from_secs(secs: u64) -> Self {
        Self {
            delay: Duration::from_secs(secs),
            attempts: Attempts::Finite(3),
        }
    }

    pub fn forever(self) -> Self {
        Self {
            attempts: Attempts::Forever,
            ..self
        }
    }

    pub async fn attempt<F: AsyncFnMut() -> Result<T, E>, T, E: core::fmt::Debug>(self, mut f: F) -> Result<T, E> {
        let mut latest= f().await;
        let mut delay = self.delay;
        'outer: loop {
            for _ in self.attempts {
                match &latest {
                    Err(e) => {
                        warn!("Operation failed: {e:?}. Retrying after {delay}");
                        Timer::after(delay).await;
                        delay *= 2;
                        latest = f().await
                    },
                    _ => break 'outer
                }
            }
        }

        latest
    }
}

#[derive(Clone, Copy)]
enum Scene {
    ParkedIdle, // Default state when booting up or waking up from long term (overnight, eg) parking, or entered when accelerometers and GPS both show zero motion for ~30 seconds
    StoplightIdle, // Entered when GPS speed is zero after decelerating
    Accelerating, // GPS speed is increasing
    Decelerating, // GPS speed is decreasing
    ParkedLongTerm, // GPS has not changed location in the last ~5 minutes
}

#[derive(Clone, Copy)]
enum Notification {
    SceneChange(Scene),
    WifiConnected,
    WifiDisconnected,
    GPSLost,
    GPSAcquired,
    BPMBeat
}

#[embassy_executor::task]
async fn render(mut events: Subscriber<'static, NoopRawMutex, Notification, 4, 4, 4>, rmt_channel: ChannelCreator<Async, 0>, gpio: GPIO5<'static>) {
    let rmt_buffer = [0u32; buffer_size_async(16)];

    let target = SmartLedsAdapterAsync::new(rmt_channel, gpio, rmt_buffer);

    // Change this number to use a different number of LEDs
    let mut pixbuf = [Default::default(); 16];

    // Change this to adjust the power available; the USB spec says 500ma is the standard limit, but sometimes you can draw more from a power brick
    const POWER_MA : u32 = 500;
    
    // You probably don't need to change these values, unless your LED strip is somehow not 5 volts
    const POWER_VOLTS : u32 = 5;
    const MAX_POWER_MW : u32 = POWER_VOLTS * POWER_MA;

    // This value is used as the 'seed' for rendering each frame, allowing us to do things like run the animation backwards, frames for double FPS, or even use system uptime for more human-paced animations
    let mut frame = 0;

    let mut target = BrightnessWriter::new(target, MAX_POWER_MW);

    info!("Rendering started!");

    loop {
        if let Some(evt) = events.try_next_message() {

        }
        let start = Instant::now();
        //pixbuf.blank();
        //output.blank().await;

        // Render the frame to the pixbuf, while also calculating the power consumption as we go
        for (coords, pix) in pixbuf.sample(&Rectangle::everything()) {
            *pix = Rgb::new(sin8(coords.x.wrapping_mul(3).wrapping_add(coords.y.wrapping_mul(3)).wrapping_add(frame)), 0, 0);
        }

        // Finally, write out the rendered frame
        //target.write(pixbuf.iter().cloned()).await.expect("Could not write frame");
        target.write(&pixbuf).await.expect("Failed to write frame!");
        //output.commit().await;

        let render_duration = Instant::now() - start;
        let render_budget = Duration::from_millis(16);

        if render_duration < render_budget {
            let remaining_budget = render_budget - render_duration;
            Timer::after(remaining_budget).await;
        } else {
            warn!("Render stall! Frame took {}ms", render_duration.as_millis());
        }

        // Increment the frame counter
        frame += 1;
    }
}