#![no_std]
#![no_main]

use core::sync::atomic::{AtomicBool, AtomicU8, Ordering};

use cortex_m_rt as _;
use kaptan_runtime::LatestImuMailbox;
use panic_halt as _;

mod board;
mod routes;

const IMU_EVENT_DRDY: u8 = 1 << 0;
const IMU_EVENT_TRANSFER_COMPLETE: u8 = 1 << 1;
const IMU_EVENT_TRANSFER_FAULT: u8 = 1 << 2;
const IMU_EVENT_TIMER: u8 = 1 << 3;
const TELEMETRY_EVENT_SAMPLE: u8 = 1 << 0;
const TELEMETRY_EVENT_HEARTBEAT: u8 = 1 << 1;
const TELEMETRY_EVENT_DMA: u8 = 1 << 2;

static IMU_EVENTS: AtomicU8 = AtomicU8::new(0);
static TELEMETRY_EVENTS: AtomicU8 = AtomicU8::new(0);
static IMU_QUEUED: AtomicBool = AtomicBool::new(false);
static TELEMETRY_QUEUED: AtomicBool = AtomicBool::new(false);
static SAMPLES: LatestImuMailbox = LatestImuMailbox::new();
static UART_COMPLETION: AtomicU8 = AtomicU8::new(0);

#[rtic::app(device = stm32f7::stm32f722, peripherals = true, dispatchers = [TIM3, TIM5])]
mod app {
    use super::{
        IMU_EVENT_DRDY, IMU_EVENT_TIMER, IMU_EVENT_TRANSFER_COMPLETE, IMU_EVENT_TRANSFER_FAULT,
        IMU_EVENTS, IMU_QUEUED, Ordering, SAMPLES, TELEMETRY_EVENT_DMA, TELEMETRY_EVENT_HEARTBEAT,
        TELEMETRY_EVENT_SAMPLE, TELEMETRY_EVENTS, TELEMETRY_QUEUED, UART_COMPLETION,
    };
    use crate::board::{
        CORE_CLOCK_MHZ, HEARTBEAT_TIMER_CONFIG, MAVLINK_COMPONENT_ID, MAVLINK_SYSTEM_ID,
        MPU6000_CONFIG, USART1_BAUD, USART1_CLOCK_HZ,
    };
    use crate::routes::KakuteMpu6000SpiRoute;
    use kaptan_dma::TxCompletion;
    use kaptan_drivers::mpu6000::Mpu6000;
    use kaptan_runtime::{ImuEvent, ImuService, TelemetryService};
    use kaptan_spi::SpiCompletion;
    use kaptan_stm32f7::stm32f722::{
        ActiveLowLedPd2, DcacheDisabledDmaMemory, DwtMicrosecondClock, Spi1, Tim2Periodic,
        Usart1Tx, Usart1TxDma, configure_exti4_pin_pa4, configure_spi1_pins_pa5_pa6_pa7,
        configure_usart1_tx_pin_pa9, select_hsi16_profile,
    };
    use stm32f7::stm32f722;

    #[shared]
    struct Shared {}

    #[local]
    struct Local {
        imu: ImuService<DwtMicrosecondClock, KakuteMpu6000SpiRoute>,
        telemetry: TelemetryService<DcacheDisabledDmaMemory, Usart1TxDma, 4>,
        led: ActiveLowLedPd2,
        heartbeat_timer: Tim2Periodic,
        heartbeat_divider: u8,
        observed_sample: u32,
    }

    #[init]
    fn init(mut context: init::Context) -> (Shared, Local) {
        let device = context.device;
        if select_hsi16_profile(&device.RCC).is_err() {
            safe_idle();
        }
        context.core.SCB.disable_dcache(&mut context.core.CPUID);
        let clock = match DwtMicrosecondClock::start::<CORE_CLOCK_MHZ>(
            &mut context.core.DCB,
            &mut context.core.DWT,
        ) {
            Ok(clock) => clock,
            Err(_) => safe_idle(),
        };

        configure_spi1_pins_pa5_pa6_pa7(&device.RCC, &device.GPIOA);
        configure_exti4_pin_pa4(&device.GPIOA);
        configure_usart1_tx_pin_pa9(&device.GPIOA);
        let spi = Spi1::new(&device.RCC, device.SPI1, device.DMA2);
        let route = KakuteMpu6000SpiRoute::new(&device.RCC, spi, device.GPIOB);
        let mut imu = ImuService::new(clock, Mpu6000::new(route, MPU6000_CONFIG));
        if imu.begin().is_err() {
            safe_idle();
        }

        let uart = match Usart1Tx::new(&device.RCC, device.USART1, USART1_CLOCK_HZ, USART1_BAUD) {
            Ok(uart) => uart,
            Err(_) => safe_idle(),
        };
        // SAFETY: D-cache is disabled above and `memory.x` excludes DTCM.
        let memory = unsafe { DcacheDisabledDmaMemory::new() };
        let telemetry = TelemetryService::new(
            MAVLINK_SYSTEM_ID,
            MAVLINK_COMPONENT_ID,
            memory,
            Usart1TxDma::new(&device.RCC, uart),
        );

        configure_exti4(&device.RCC, device.SYSCFG, device.EXTI);
        (
            Shared {},
            Local {
                imu,
                telemetry,
                led: ActiveLowLedPd2::new(&device.RCC, device.GPIOD),
                heartbeat_timer: Tim2Periodic::new(
                    &device.RCC,
                    device.TIM2,
                    HEARTBEAT_TIMER_CONFIG,
                ),
                heartbeat_divider: 0,
                observed_sample: 0,
            },
        )
    }

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

    #[task(binds = EXTI4, priority = 15)]
    fn exti4(_: exti4::Context) {
        // SAFETY: this handler exclusively acknowledges the board's EXTI4 bit.
        unsafe {
            (&*stm32f722::EXTI::ptr())
                .pr()
                .write(|write| write.pr4().clear())
        };
        schedule_imu(IMU_EVENT_DRDY);
    }

    #[task(binds = DMA2_STREAM0, priority = 14)]
    fn spi_rx_dma(_: spi_rx_dma::Context) {
        match Spi1::acknowledge_rx_interrupt() {
            SpiCompletion::Complete => schedule_imu(IMU_EVENT_TRANSFER_COMPLETE),
            SpiCompletion::TransferFault => schedule_imu(IMU_EVENT_TRANSFER_FAULT),
            SpiCompletion::NotPending => {}
        }
    }

    #[task(binds = DMA2_STREAM3, priority = 14)]
    fn spi_tx_dma(_: spi_tx_dma::Context) {
        if Spi1::acknowledge_tx_interrupt() == SpiCompletion::TransferFault {
            schedule_imu(IMU_EVENT_TRANSFER_FAULT);
        }
    }

    #[task(binds = TIM2, priority = 4, local = [heartbeat_timer])]
    fn tim2(context: tim2::Context) {
        context.local.heartbeat_timer.acknowledge_update();
        schedule_imu(IMU_EVENT_TIMER);
        schedule_telemetry(TELEMETRY_EVENT_HEARTBEAT);
    }

    #[task(binds = DMA2_STREAM7, priority = 10)]
    fn uart_dma(_: uart_dma::Context) {
        match Usart1TxDma::acknowledge_interrupt() {
            TxCompletion::Complete => {
                UART_COMPLETION.store(1, Ordering::Release);
                schedule_telemetry(TELEMETRY_EVENT_DMA);
            }
            TxCompletion::TransferFault | TxCompletion::Halted => {
                UART_COMPLETION.store(2, Ordering::Release);
                schedule_telemetry(TELEMETRY_EVENT_DMA);
            }
            TxCompletion::NotPending => {}
        }
    }

    #[task(priority = 13, local = [imu])]
    async fn imu_worker(context: imu_worker::Context) {
        let events = IMU_EVENTS.swap(0, Ordering::AcqRel);
        if events & IMU_EVENT_TRANSFER_FAULT != 0 {
            if !matches!(
                context
                    .local
                    .imu
                    .on_transfer_complete(SpiCompletion::TransferFault),
                ImuEvent::Fault(_)
            ) {
                panic!("SPI fault without an active MPU transaction");
            }
            panic!("SPI DMA transfer fault");
        }
        if events & IMU_EVENT_TRANSFER_COMPLETE != 0 {
            match context
                .local
                .imu
                .on_transfer_complete(SpiCompletion::Complete)
            {
                ImuEvent::Sample(sample) => {
                    SAMPLES.publish(sample);
                    schedule_telemetry(TELEMETRY_EVENT_SAMPLE);
                }
                ImuEvent::Fault(_) => panic!("MPU transfer failed"),
                ImuEvent::Initialized | ImuEvent::Idle => {}
            }
        }
        if events & IMU_EVENT_TIMER != 0
            && matches!(context.local.imu.on_periodic_tick(), ImuEvent::Fault(_))
        {
            panic!("MPU initialization failed");
        }
        if events & IMU_EVENT_DRDY != 0
            && matches!(context.local.imu.on_data_ready(), ImuEvent::Fault(_))
        {
            panic!("MPU sample submission failed");
        }
        IMU_QUEUED.store(false, Ordering::Release);
        if IMU_EVENTS.load(Ordering::Acquire) != 0 && !IMU_QUEUED.swap(true, Ordering::AcqRel) {
            if imu_worker::spawn().is_err() {
                panic!("IMU worker reschedule failed");
            }
        }
    }

    #[task(priority = 9, local = [telemetry, led, heartbeat_divider, observed_sample])]
    async fn telemetry_worker(context: telemetry_worker::Context) {
        let events = TELEMETRY_EVENTS.swap(0, Ordering::AcqRel);
        if events & TELEMETRY_EVENT_DMA != 0 {
            let completion = match UART_COMPLETION.swap(0, Ordering::AcqRel) {
                1 => TxCompletion::Complete,
                2 => TxCompletion::TransferFault,
                _ => TxCompletion::NotPending,
            };
            if context.local.telemetry.on_dma_complete(completion).is_err() {
                panic!("telemetry DMA failed");
            }
        }
        if events & TELEMETRY_EVENT_SAMPLE != 0 {
            if let Some(sample) = SAMPLES.take_newer(context.local.observed_sample) {
                if context.local.telemetry.publish_sample(&sample).is_err() {
                    panic!("RAW_IMU publication failed");
                }
            }
        }
        if events & TELEMETRY_EVENT_HEARTBEAT != 0 {
            *context.local.heartbeat_divider = context.local.heartbeat_divider.wrapping_add(1);
            if *context.local.heartbeat_divider == 10 {
                *context.local.heartbeat_divider = 0;
                if context.local.telemetry.publish_heartbeat().is_err() {
                    panic!("heartbeat publication failed");
                }
                context.local.led.toggle();
            }
        }
        TELEMETRY_QUEUED.store(false, Ordering::Release);
        if TELEMETRY_EVENTS.load(Ordering::Acquire) != 0
            && !TELEMETRY_QUEUED.swap(true, Ordering::AcqRel)
        {
            if telemetry_worker::spawn().is_err() {
                panic!("telemetry worker reschedule failed");
            }
        }
    }

    fn schedule_imu(event: u8) {
        IMU_EVENTS.fetch_or(event, Ordering::Release);
        if !IMU_QUEUED.swap(true, Ordering::AcqRel) && imu_worker::spawn().is_err() {
            panic!("IMU worker schedule failed");
        }
    }

    fn schedule_telemetry(event: u8) {
        TELEMETRY_EVENTS.fetch_or(event, Ordering::Release);
        if !TELEMETRY_QUEUED.swap(true, Ordering::AcqRel) && telemetry_worker::spawn().is_err() {
            panic!("telemetry worker schedule failed");
        }
    }

    fn configure_exti4(rcc: &stm32f722::RCC, syscfg: stm32f722::SYSCFG, exti: stm32f722::EXTI) {
        rcc.apb2enr().modify(|_, write| write.syscfgen().set_bit());
        syscfg
            .exticr2()
            .modify(|_, write| unsafe { write.exti4().bits(0) });
        exti.imr().modify(|_, write| write.im4().set_bit());
        exti.rtsr().modify(|_, write| write.tr4().set_bit());
        exti.ftsr().modify(|_, write| write.tr4().clear_bit());
        exti.pr().write(|write| write.pr4().clear());
    }

    fn safe_idle() -> ! {
        loop {
            cortex_m::asm::wfi();
        }
    }
}