#![no_std]
#![no_main]

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

use common::{Command, Response, SensorData, BAUDRATE};
use embassy_executor::Spawner;
use embassy_rp::{bind_interrupts, block::ImageDef, gpio::{Level, Output}, peripherals::{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};
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;
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>;
});

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

#[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();

    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 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; 1024]> = ConstStaticCell::new([0u8;1024]);
    let tx_buf = TX_BUF.take();
    static RX_BUF: ConstStaticCell<[u8; 1024]> = ConstStaticCell::new([0u8;1024]);
    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 STARTUP: OnceLock<()> = OnceLock::new();

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

    let enabled = &ENABLED;
    let mut enable_watchdog = Instant::now();
    let enable_watchdog_time = Duration::from_millis(150);

    info!("ready");

    loop {
        let command = if enabled.load(Ordering::Acquire) { 
            let Ok(command) = with_deadline(enable_watchdog + enable_watchdog_time, COMMANDS.receive()).await else {
                warn!("no message received");
                enabled.store(false, Ordering::Release);
                continue;
            };
            command
        } else {
            // stop all motors
            drive_pwm.set_config(&stopped);
            info!("waiting for command");
            COMMANDS.receive().await
        };

        STARTUP.get_or_init(||());

        match command {
            Command::Twist(forward, right) => {
                drive_conf.compare_a = calc_speed(-forward - right);
                drive_conf.compare_b = calc_speed(forward - right);
                drive_pwm.set_config(&drive_conf);
            },
            Command::Stop => { 
                drive_pwm.set_config(&stopped);
            },
            Command::Enable => {
                enabled.store(true, Ordering::Release);
                enable_watchdog = Instant::now();
            },
            Command::Disable => {
                enabled.store(false, Ordering::Release);
                drive_pwm.set_config(&stopped);
            },
            Command::FeedWatchdog => {
                enable_watchdog = Instant::now();
            }
            c => { error!("{c:?} unimplemented") }
        };

        info!("Blink");
    }
}

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

    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];
        rx.read_exact(&mut buf[..len]).await.unwrap();

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

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

        COMMANDS.send(data).await;
    }
}

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

        let packet = Response {
            enabled: enabled.load(Ordering::Acquire),
            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:?}");
        }
        tx.flush().await.unwrap();
    }
}

/// -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
}