#![no_std]
#![no_main]

use core::{panic::PanicInfo, sync::atomic::{AtomicU16, Ordering}};

use common::{Response, RobotState, SensorData, BAUDRATE};
use embassy_executor::Spawner;
use embassy_rp::{adc::{self, Adc}, bind_interrupts, block::ImageDef, gpio::{Level, Output, Pull}, peripherals::{ADC, ADC_TEMP_SENSOR, PIN_28, UART0, UART1, USB}, pwm::{self, Pwm}, uart::{BufferedInterruptHandler, BufferedUart, BufferedUartRx, BufferedUartTx, Config}, usb::Driver};
use embassy_sync::{blocking_mutex::raw::CriticalSectionRawMutex, channel::Channel, once_lock::OnceLock, watch::Watch};
use embassy_time::{with_deadline, with_timeout, Duration, Instant, Timer};
use embedded_io_async::{BufRead, Read, Write};
use log::{error, info, trace, warn};
use nalgebra::{clamp, Vector3};
use portable_atomic::AtomicBool;
use static_cell::{ConstStaticCell, StaticCell};

bind_interrupts!(struct Irqs {
    I2C1_IRQ => embassy_rp::i2c::InterruptHandler<embassy_rp::peripherals::I2C1>;
    USBCTRL_IRQ => embassy_rp::usb::InterruptHandler<USB>;
    UART1_IRQ => BufferedInterruptHandler<UART1>;
    ADC_IRQ_FIFO => adc::InterruptHandler;
});

pub static COMMANDS: Watch<CriticalSectionRawMutex, (RobotState, Instant), 5> = Watch::new();
pub static SENSOR_DATA: Channel<CriticalSectionRawMutex, SensorData, 5> = Channel::new();

pub static BUS_VOLTAGE: AtomicU16 = AtomicU16::new(u16::MAX);
const BUS_ADC_TO_VOLTS: f32 = 251.6763848397;

/// volts
fn bus_voltage() -> f32 {
    BUS_VOLTAGE.load(Ordering::Acquire) as f32 / BUS_ADC_TO_VOLTS
}

#[link_section = ".start_block"]
#[used]
pub static IMAGE_DEF: ImageDef = ImageDef::secure_exe();

#[embassy_executor::task]
async fn logger_task(driver: Driver<'static, USB>) {
    embassy_usb_logger::run!(1024, log::LevelFilter::Trace, driver);
}

#[panic_handler]
fn panic( info: &PanicInfo) -> ! {
    error!("{}", info);
    loop { }
}

#[embassy_executor::main]
async fn main(spawner: Spawner) {
    let p = embassy_rp::init(Default::default());

    let mut led = Output::new(p.PIN_25, Level::Low);

    led.set_high();

    let driver = Driver::new(p.USB, Irqs);
    spawner.spawn(logger_task(driver)).unwrap();
    info!("logger enabled");

    let mut drive_conf: pwm::Config = Default::default();
    drive_conf.divider = 150.into();
    drive_conf.top = 20000; // 20ms
    drive_conf.compare_b = 1500; // 1.5ms
    drive_conf.compare_a = 1500; // 1.5ms
    let stopped = drive_conf.clone();
    let mut drive_pwm = Pwm::new_output_ab(p.PWM_SLICE0, p.PIN_16, p.PIN_17, stopped.clone());

    let disabled_orange = Vector3::new(254, 100, 0);
    let enabled_green = Vector3::new(0, 140, 0);

    let mut light_conf_rg: pwm::Config = Default::default();
    light_conf_rg.divider = 1.into();
    light_conf_rg.top = 255;
    light_conf_rg.compare_b = 0;
    light_conf_rg.compare_a = 0;
    let mut light_pwm_rg = Pwm::new_output_ab(p.PWM_SLICE7, p.PIN_14, p.PIN_15, light_conf_rg);

    let mut light_conf_b: pwm::Config = Default::default();
    light_conf_b.divider = 1.into();
    light_conf_b.top = 255;
    light_conf_b.compare_b = 0;
    let mut light_pwm_b = Pwm::new_output_b(p.PWM_SLICE6, p.PIN_13, light_conf_b);
    set_underglow(&mut light_pwm_rg, &mut light_pwm_b, disabled_orange.clone());

    Timer::after_secs(13).await; // pi 0 serial ports act strange during the boot process

    spawner.spawn(bus_voltage_monitor(p.ADC, p.PIN_28, p.ADC_TEMP_SENSOR)).unwrap();

    let config = embassy_rp::i2c::Config::default();
    let bus = embassy_rp::i2c::I2c::new_async(p.I2C1, p.PIN_19, p.PIN_18, Irqs, config);

    static TX_BUF: ConstStaticCell<[u8; 2048]> = ConstStaticCell::new([0u8;2048]);
    let tx_buf = TX_BUF.take();
    static RX_BUF: ConstStaticCell<[u8; 2048]> = ConstStaticCell::new([0u8;2048]);
    let rx_buf = RX_BUF.take();

    let mut uart_config = Config::default();
    uart_config.baudrate = BAUDRATE;
    let uart = BufferedUart::new(p.UART1, Irqs, p.PIN_20, p.PIN_21, tx_buf, rx_buf, uart_config);

    let (tx,rx) = uart.split();

    spawner.spawn(decoder(rx)).unwrap();

    static ENABLED: AtomicBool = AtomicBool::new(false);
    spawner.spawn(telemetry(tx, &ENABLED)).unwrap();

    let watchdog_time = Duration::from_millis(150);

    info!("ready");

    let mut commands = COMMANDS.receiver().unwrap();

    loop {
        let motor_update = Timer::after_millis(3);
        let (command, time) = commands.get().await;

        if time.elapsed() > watchdog_time {
            drive_pwm.set_config(&stopped);
            set_underglow(&mut light_pwm_rg, &mut light_pwm_b, disabled_orange.clone());
            ENABLED.store(false, Ordering::Acquire);
            commands.changed().await;
            set_underglow(&mut light_pwm_rg, &mut light_pwm_b, enabled_green.clone());
            ENABLED.store(true, Ordering::Acquire);
            continue;
        }

        ENABLED.store(true, Ordering::Acquire);


        // 100% effort at 9.5v, 0% effort at 7.5v
        let brownout = ((bus_voltage()-7.5)/2.0).clamp(0., 1.);
        let command = command.brownout(brownout);

        // drive motors
        let (forward, right) = command.drive;
        drive_conf.compare_a = calc_speed(-forward - right);
        drive_conf.compare_b = calc_speed(forward - right);
        drive_pwm.set_config(&drive_conf);

        // underglow
        set_underglow(&mut light_pwm_rg, &mut light_pwm_b, command.underglow);

        motor_update.await;
    }
}

/// Receive data from the pi 0 over UART and deserialize it
#[embassy_executor::task]
async fn decoder(mut rx: BufferedUartRx<'static, UART1>) {
    info!("Reading...");
    let sender = COMMANDS.sender();

    loop {
        let mut len = [0;4];
        rx.read_exact(&mut len).await.unwrap();
        let len = u32::from_be_bytes(len) as usize;

        let mut buf = [0;1024];
        if let Err(e) = rx.read_exact(&mut buf[..len]).await {
            trace!("read fail {e:?}");
        };

        let Ok(data) = postcard::from_bytes::<RobotState>(&buf[..len]) else {
            error!("message decode fail");
            continue;
        };

        trace!("received {data:?}");

        sender.send((data, Instant::now()));
    }
}

/// Receive data from channel and send it over UART
#[embassy_executor::task]
async fn telemetry(mut tx: BufferedUartTx<'static, UART1>, enabled: &'static AtomicBool) {
    info!("Telemetry waiting...");
    //COMMANDS.receiver().unwrap().get().await;
    info!("Telemetry started...");
    loop {
        let data = with_timeout(Duration::from_millis(20), SENSOR_DATA.receive()).await.ok();

        if !enabled.load(Ordering::Release) {
            Timer::after_millis(10).await;
            continue;
        }

        let packet = Response {
            sensor_data: data,
            uptime_micros: Instant::now().as_micros(),
        };

        let Ok(serialized) = postcard::to_vec_cobs::<_, 1024>(&packet) else {
            error!("serialization error");
            continue;
        };

        if let Err(e) = tx.write_all(&serialized).await {
            error!("transport error {e:?}");
        }
        if let Err(e) = tx.flush().await {
            error!("flush err {e:?}");
        }
    }
}

/// -1 to 1
fn calc_speed(speed: f32) -> u16 {
    const COUNTS_PER_MS: f32 = 1000.;
    let speed = speed.clamp(-1., 1.);

    let ms = (speed/2.0)+1.5;

    (COUNTS_PER_MS * ms) as u16
}

#[embassy_executor::task]
async fn bus_voltage_monitor(adc: ADC, bus: PIN_28, temp: ADC_TEMP_SENSOR) {
    let mut adc = Adc::new(adc, Irqs, adc::Config::default());
    let mut bus_voltage = adc::Channel::new_pin(bus, Pull::None);
    let mut ts = adc::Channel::new_temp_sensor(temp);

    // from embassy examples
    fn convert_to_celsius(raw_temp: u16) -> f32 {
        let temp = 27.0 - (raw_temp as f32 * 3.3 / 4096.0 - 0.706) / 0.001721;
        let sign = if temp < 0.0 { -1.0 } else { 1.0 };
        let rounded_temp_x10: i16 = ((temp * 10.0) + 0.5 * sign) as i16;
        (rounded_temp_x10 as f32) / 10.0
    }

    loop {
        let level = adc.read(&mut bus_voltage).await.unwrap();
        BUS_VOLTAGE.store(level, Ordering::Release);
        // empirically calculated against $20 microcenter voltmeter (10k & 33k divider circuit)
        SENSOR_DATA.send(SensorData::BusVoltage(level as f32 / BUS_ADC_TO_VOLTS)).await;
        let temp = adc.read(&mut ts).await.unwrap();
        SENSOR_DATA.send(SensorData::AmbientTemperature(convert_to_celsius(temp))).await;
        Timer::after_millis(3).await;
    }
}

fn set_underglow(rg: &mut Pwm, b: &mut Pwm, color: Vector3<u8>) {
    let mut light_conf_rg: pwm::Config = Default::default();
    light_conf_rg.divider = 1.into();
    light_conf_rg.top = 255;
    light_conf_rg.compare_b = color[1] as u16;
    light_conf_rg.compare_a = color[0] as u16;
    rg.set_config(&light_conf_rg);

    let mut light_conf_b: pwm::Config = Default::default();
    light_conf_b.divider = 1.into();
    light_conf_b.top = 255;
    light_conf_b.compare_b = color[2] as u16;
    b.set_config(&light_conf_b);
}