#![no_std]
#![no_main]
#![feature(abi_avr_interrupt)]

use core::cell::RefCell;
use core::f32;
use core::mem::MaybeUninit;

use arduino_hal::{clock::MHz16, hal::usart::UsartWriter};
use arduino_hal::hal::port::*;
use arduino_hal::hal::usart::UsartReader;
use arduino_hal::pac::USART0;
use arduino_hal::port::mode::*;
use arduino_hal::prelude::*;
use avr_device::interrupt::{self, Mutex};
use common::{Command, Telemetry};
use panic_halt as _;

fn shift_out(data_pin: &mut Pin<Output, PB0>, clock_pin: &mut Pin<Output, PD4>, data: &u8) {
    for i in 0..8 {
        let n = data & (1 << i);

        if n == 0 {
            data_pin.set_low();
        } else {
            data_pin.set_high();
        }

        clock_pin.set_high();
        clock_pin.set_low();
    }
}

fn update_shift_register(
    data_pin: &mut Pin<Output, PB0>,
    latch_pin: &mut Pin<Output, PB4>,
    clock_pin: &mut Pin<Output, PD4>,
    data: &u8,
) {
    latch_pin.set_low();

    shift_out(data_pin, clock_pin, data);

    latch_pin.set_high();
}

static mut UART_RX: MaybeUninit<UsartReader<USART0, Pin<Input, PD0>, Pin<Output, PD1>, MHz16>> = MaybeUninit::uninit();

static INPUT_LINE: Mutex<RefCell<[u8; 20]>> = Mutex::new(RefCell::new([0;20]));

#[avr_device::interrupt(atmega328p)]
unsafe fn USART_RX() {
    let rx = &mut *UART_RX.as_mut_ptr();

    static mut BUF: [u8; 20] = [0;20];
    static mut N: usize = 0; // index into line

    if let Ok(val) = rx.read() {
        if val == b'\0' {

            interrupt::free(|cs| {
                let mut line = INPUT_LINE.borrow(cs).borrow_mut();
                *line = BUF;
                line[N] = b'\0';
            });

            N = 0;
        } else {
            BUF[N] = val;
            N += 1;
            if N == 20 {
                N = 0;
            }
        }
    }
}


#[arduino_hal::entry]
fn main() -> ! {
    let dp = arduino_hal::Peripherals::take().unwrap();
    let pins = arduino_hal::pins!(dp);

    // shift register
    let mut data_pin = pins.d8.into_output();
    let mut clock_pin = pins.d4.into_output();
    let mut enable_pin = pins.d7.into_output();
    let mut latch_pin = pins.d12.into_output();

    enable_pin.set_low(); // enable outputs

    pins.d3.into_output(); // right
    pins.d11.into_output(); // left

    // TCCR2A |= _BV(COM2A1) | _BV(WGM20) | _BV(WGM21); // fast PWM, turn on oc2a
    // TCCR2B = freq & 0x7;
    // OCR2A = 0;
    let tc2 = dp.TC2;
    tc2.tccr2a.write(|w| { w
            .wgm2().pwm_fast()
            .com2a().match_clear()
            .com2b().match_clear()
    });
    tc2.tccr2b.write(|w| w.cs2().prescale_8().wgm22().clear_bit());
    tc2.ocr2a.write(|w| w.bits(0));
    tc2.ocr2b.write(|w| w.bits(0));

    let mut trig = pins.a4.into_output();
    let echo = pins.a5;

    let ultrasonic_timer = dp.TC1;
    // 4us per tick
    ultrasonic_timer.tccr1b.write(|w| w.cs1().prescale_64());

    let mut led = pins.d13.into_output();

    let mut serial = arduino_hal::default_serial!(dp, pins, 115200);

    avr_device::interrupt::disable();
    serial.listen(arduino_hal::hal::usart::Event::RxComplete);

    let (rx,mut tx) = serial.split();

    unsafe { 
        UART_RX = MaybeUninit::new(rx);
        avr_device::interrupt::enable();
    };

    let mut telem = Telemetry { distance: f32::NAN };

    loop {
        let mut shift_register = 0;
        const LEFT_FWD: u8 = 1 << 5; // 1a
        const LEFT_REV: u8 = 1 << 4; // 1b
        const RIGHT_FWD: u8 = 1 << 3; // 2b
        const RIGHT_REV: u8 = 1 << 6; // 2a

        led.toggle();

        let command = decode_command();

        if command.left > 0. {
            shift_register |= LEFT_FWD;
        }
        if command.left < 0. {
            shift_register |= LEFT_REV;
        }

        if command.right > 0. {
            shift_register |= RIGHT_FWD;
        }
        if command.right < 0. {
            shift_register |= RIGHT_REV;
        }

        fn to_pwm(val: f32) -> u8 {
            (val.abs().min(1.0) * 255.0) as u8
        }

        // 16/255
        tc2.ocr2a.write(|w| w.bits(to_pwm(command.left))); // left
        tc2.ocr2b.write(|w| w.bits(to_pwm(command.right))); // right

        update_shift_register(&mut data_pin, &mut latch_pin, &mut clock_pin, &shift_register);

        // reset timer
        ultrasonic_timer.tcnt1.write(|w| w.bits(0));

        // send ultrasonic pulse
        trig.set_high();
        arduino_hal::delay_us(10);
        trig.set_low();

        let mut response = true;

        /// 0.1s/4µs = 25,000
        const SIGNAL_START_TIMEOUT: u16 = 25000;

        // wait for return signal start
        while echo.is_low() {
            if ultrasonic_timer.tcnt1.read().bits() >= SIGNAL_START_TIMEOUT {
                response = false;
                break;
            }
        }

        if response {
            let remainder = ultrasonic_timer.tcnt1.read().bits();
            ultrasonic_timer.tcnt1.write(|w| w.bits(0));

            // Wait for the echo to get low again
            while echo.is_high() {}
            let counts = ultrasonic_timer.tcnt1.read().bits();
            let microseconds = counts.saturating_mul(4);

            //  inch per 148 microseconds
            telem.distance = microseconds as f32 / 148.;
            if microseconds == u16::MAX { telem.distance = f32::NAN };

            // delay for the rest of the return signal timeout if any is left
            ultrasonic_timer.tcnt1.write(|w| w.bits(SIGNAL_START_TIMEOUT.saturating_sub(remainder + counts)));
            while ultrasonic_timer.tcnt1.read().bits() <= SIGNAL_START_TIMEOUT {}

        } else {
            telem.distance = f32::NAN;
        }

        let mut buf = [0;20];
        let encoded = postcard::to_slice_cobs(&telem, &mut buf).unwrap();

        for word in encoded {
            let _ = tx.write(*word);
            let _ = tx.flush();
            arduino_hal::delay_us(100);
        }

        arduino_hal::delay_ms(250);
    }
}

fn decode_command() -> Command {
    let mut input = interrupt::free(|cs| {
        INPUT_LINE.borrow(cs).borrow().clone()
    });

    let Some(length) = input.iter().position(|n| *n == 0) else { return Default::default() };

    postcard::from_bytes_cobs(&mut input[..=length]).unwrap_or(Default::default())
}