#![no_std]
#![no_main]

use embedded_alloc::Heap;
use panic_halt as _;

use fugit::RateExtU32;

use cortex_m::prelude::_embedded_hal_watchdog_Watchdog;
use embedded_hal::adc::OneShot;
use hal::i2c;
use ssd1306::{prelude::*, I2CDisplayInterface, Ssd1306};
pub use waveshare_rp2040_zero::hal;

use embedded_graphics::text::Text;
use embedded_graphics::{
    mono_font::{ascii::FONT_10X20, MonoTextStyle},
    pixelcolor::BinaryColor,
    prelude::*,
    text::Alignment,
};
use smart_leds::{SmartLedsWrite, RGB8};

extern crate alloc;
use alloc::format;

use fugit::MicrosDurationU32;
use hal::{
    gpio::{
        bank0::{Gpio14, Gpio15, Gpio16},
        Pin,
    },
    pac::I2C1,
    prelude::*,
    timer::{Alarm, Alarm0, Alarm1},
    Adc,
};

#[global_allocator]
static HEAP: Heap = Heap::empty();

const HEAP_SIZE: usize = 1024;

const XTAL_FREQ_HZ: u32 = 12_000_000u32;
const MVG_AVG_COUNT: u16 = 20;

#[rtic::app(device = hal::pac, peripherals = true)]
mod app {

    const DISPLAY_UPDATE_INT_TICKS: u32 = 500_000;
    const SENSOR_READ_INT_TICKS: u32 = 50_000;

    use super::*;
    type DisplayI2C = hal::I2C<
        I2C1,
        (
            Pin<Gpio14, hal::gpio::Function<hal::gpio::I2C>>,
            Pin<Gpio15, hal::gpio::Function<hal::gpio::I2C>>,
        ),
    >;
    type Display = Ssd1306<
        I2CInterface<DisplayI2C>,
        DisplaySize128x64,
        ssd1306::mode::BufferedGraphicsMode<DisplaySize128x64>,
    >;
    type LED =
        ws2812_pio::Ws2812<hal::pac::PIO0, hal::pio::SM0, hal::timer::CountDown<'static>, Gpio16>;

    #[shared]
    struct Shared {
        alarm0: Alarm0,
        alarm1: Alarm1,
        sensor_value: u16,
        led: LED,
    }

    #[local]
    struct Local {
        display: Display,
        adc: Adc,
        adc_pin: Pin<hal::gpio::bank0::Gpio29, hal::gpio::FloatingInput>,
        watchdog: hal::Watchdog,
    }

    #[init()]
    fn init(c: init::Context) -> (Shared, Local, init::Monotonics) {
        init_heap();

        let mut resets = c.device.RESETS;
        let sio = hal::Sio::new(c.device.SIO);
        let mut watchdog = hal::Watchdog::new(c.device.WATCHDOG);

        let clocks = hal::clocks::init_clocks_and_plls(
            XTAL_FREQ_HZ,
            c.device.XOSC,
            c.device.CLOCKS,
            c.device.PLL_SYS,
            c.device.PLL_USB,
            &mut resets,
            &mut watchdog,
        )
        .ok()
        .unwrap();

        let pins = rp2040_hal::gpio::Pins::new(
            c.device.IO_BANK0,
            c.device.PADS_BANK0,
            sio.gpio_bank0,
            &mut resets,
        );

        let mut timer = hal::Timer::new(c.device.TIMER, &mut resets);
        let mut alarm0 = timer.alarm_0().unwrap();
        let _ = alarm0.schedule(MicrosDurationU32::from_ticks(DISPLAY_UPDATE_INT_TICKS));
        alarm0.enable_interrupt();

        let mut alarm1 = timer.alarm_1().unwrap();
        let _ = alarm1.schedule(MicrosDurationU32::from_ticks(SENSOR_READ_INT_TICKS));
        alarm1.enable_interrupt();

        let (mut pio, sm0, _, _, _) = c.device.PIO0.split(&mut resets);
        let mut led = ws2812_pio::Ws2812::new(
            pins.gpio16.into_mode(),
            &mut pio,
            sm0,
            clocks.peripheral_clock.freq(),
            unsafe { core::mem::transmute(timer.count_down()) },
        );

        led.write([RGB8::new(255, 255, 0)].iter().copied()).unwrap();

        let scl = pins.gpio15.into_mode();
        let sda = pins.gpio14.into_mode();
        let i2c = i2c::I2C::i2c1(
            c.device.I2C1,
            sda,
            scl,
            100.kHz(),
            &mut resets,
            125_000_000.Hz(),
        );
        let display = init_display(i2c);

        let mut adc = Adc::new(c.device.ADC, &mut resets);
        let mut adc_pin = pins.gpio29.into_floating_input();
        led.write([RGB8::new(0, 255, 255)].iter().cloned()).unwrap();

        (
            Shared {
                alarm0,
                alarm1,
                sensor_value: adc.read(&mut adc_pin).unwrap(),
                led,
            },
            Local {
                display,
                adc,
                adc_pin,
                watchdog,
            },
            init::Monotonics(),
        )
    }

    fn init_heap() {
        use core::mem::MaybeUninit;
        static mut HEAP_MEM: [MaybeUninit<u8>; HEAP_SIZE] = [MaybeUninit::uninit(); HEAP_SIZE];
        unsafe { HEAP.init(HEAP_MEM.as_ptr() as usize, HEAP_SIZE) }
    }

    fn init_display(i2c: DisplayI2C) -> Display {
        let interface = I2CDisplayInterface::new(i2c);
        let mut display = Ssd1306::new(interface, DisplaySize128x64, DisplayRotation::Rotate0)
            .into_buffered_graphics_mode();
        display.init().unwrap();

        Text::with_alignment(
            "Initialized",
            display.bounding_box().center(),
            MonoTextStyle::new(&FONT_10X20, BinaryColor::On),
            Alignment::Center,
        )
        .draw(&mut display)
        .unwrap();

        display.flush().unwrap();

        display
    }

    #[task(binds = TIMER_IRQ_0, priority = 1, shared = [alarm0, sensor_value], local = [display])]
    fn update_display(mut c: update_display::Context) {
        let text = c.shared.sensor_value.lock(|val| format!("{val}"));

        let display = c.local.display;
        display.clear();
        Text::with_alignment(
            &text,
            display.bounding_box().center(),
            MonoTextStyle::new(&FONT_10X20, BinaryColor::On),
            Alignment::Center,
        )
        .draw(display)
        .unwrap();
        display.flush().unwrap();

        c.shared.alarm0.lock(|a| {
            a.clear_interrupt();
            let _ = a.schedule(MicrosDurationU32::from_ticks(DISPLAY_UPDATE_INT_TICKS));
        });
    }

    #[task(binds = TIMER_IRQ_1, priority = 2, shared = [alarm1, sensor_value, led], local = [adc, adc_pin, watchdog])]
    fn read_sensor(mut c: read_sensor::Context) {
        let mut sensor_value: u16 = c.local.adc.read(c.local.adc_pin).unwrap();
        sensor_value = c.shared.sensor_value.lock(|v| {
            *v = (*v * (MVG_AVG_COUNT - 1) + sensor_value) / MVG_AVG_COUNT;
            *v
        });
        let color = match sensor_value {
            0..=450 => RGB8::new(255, 0, 0),
            451.. => RGB8::new(0, 255, 0),
        };

        c.shared.led.lock(|l| {
            l.write([color].iter().cloned()).unwrap();
        });

        c.local.watchdog.feed();
        c.shared.alarm1.lock(|a| {
            a.clear_interrupt();
            let _ = a.schedule(MicrosDurationU32::from_ticks(DISPLAY_UPDATE_INT_TICKS));
        });
    }

    #[idle]
    fn idle(_cx: idle::Context) -> ! {
        loop {
            cortex_m::asm::nop();
        }
    }
}