bioz-firmware-rs/src/communication.rs

use defmt::info;

use embassy_executor::Spawner;
use embassy_sync::blocking_mutex::raw::{ThreadModeRawMutex, CriticalSectionRawMutex};
use embassy_sync::signal::Signal;
use embassy_futures::select::{select, Either};
use embassy_stm32::usb::Driver;
use embassy_stm32::{peripherals, uid, usb};

use heapless::Vec;
use postcard_rpc::{
    define_dispatch,
    header::VarHeader,
    server::{
        impls::embassy_usb_v0_4::{
            dispatch_impl::{spawn_fn, WireRxBuf, WireRxImpl, WireSpawnImpl, WireStorage, WireTxImpl},
            PacketBuffers,
        },
        Dispatch, Server, Sender, SpawnContext,
    },
};

use bioz_icd_rs::{GetUniqueIdEndpoint, MeasurementPointSet, MultiImpedanceOutputTopic, MultiImpedanceResult, MultiImpedanceStartRequest, PingEndpoint, SetGreenLedEndpoint, SingleImpedanceOutputTopic, SingleImpedanceStartRequest, StartMultiImpedanceEndpoint, StartSingleImpedanceEndpoint, StopMultiImpedanceEndpoint, StopSingleImpedanceEndpoint, ENDPOINT_LIST, TOPICS_IN_LIST, TOPICS_OUT_LIST};

use crate::impedance::{ImpedanceSetupType, RunningMode, IMPEDANCE_CHANNEL_MULTI, IMPEDANCE_CHANNEL_SINGLE};

// Postcard RPC types
type AppDriver = usb::Driver<'static, peripherals::USB>;
type AppStorage = WireStorage<ThreadModeRawMutex, AppDriver, 256, 256, 64, 256>;
type BufStorage = PacketBuffers<1024, 1024>;
type AppTx = WireTxImpl<ThreadModeRawMutex, AppDriver>;
type AppRx = WireRxImpl<AppDriver>;
type AppServer = Server<AppTx, AppRx, WireRxBuf, MyApp>;

use static_cell::ConstStaticCell;
static PBUFS: ConstStaticCell<BufStorage> = ConstStaticCell::new(BufStorage::new());
static STORAGE: AppStorage = AppStorage::new();

pub struct Context {
    pub unique_id: [u8; 12],
    pub impedance_setup: &'static ImpedanceSetupType,
}

pub struct SpawnCtx {
    pub impedance_setup: &'static ImpedanceSetupType,
}

impl SpawnContext for Context {
    type SpawnCtxt = SpawnCtx;
    fn spawn_ctxt(&mut self) -> Self::SpawnCtxt {
        SpawnCtx {
            impedance_setup: self.impedance_setup,
        }
    }
}

define_dispatch! {
    app: MyApp;
    spawn_fn: spawn_fn;
    tx_impl: AppTx;
    spawn_impl: WireSpawnImpl;
    context: Context;

    endpoints: {
        list: ENDPOINT_LIST;

        | EndpointTy                    | kind      | handler                        |
        | ----------                    | ----      | -------                        |
        | PingEndpoint                  | blocking  | ping_handler                   |
        | GetUniqueIdEndpoint           | blocking  | get_unique_id_handler          |
        | SetGreenLedEndpoint           | async     | set_green_led_handler          |
        | StartSingleImpedanceEndpoint  | spawn     | start_single_impedance_handler |
        | StopSingleImpedanceEndpoint   | async     | stop_single_impedance_handler  |
        | StartMultiImpedanceEndpoint   | spawn     | start_multi_impedance_handler   |
        | StopMultiImpedanceEndpoint    | async     | stop_multi_impedance_handler    |
    };
    topics_in: {
        list: TOPICS_IN_LIST;

        | TopicTy                   | kind      | handler                       |
        | ----------                | ----      | -------                       |
    };
    topics_out: {
        list: TOPICS_OUT_LIST;
    };
}

fn usb_config() -> embassy_usb::Config<'static> {
    let mut config = embassy_usb::Config::new(0x16c0, 0x27DD);
    config.manufacturer = Some("Hubald Verzijl");
    config.product = Some("Bioz Amplifier");
    config.serial_number = Some("12345678");

    // Required for windows compatibility.
    // https://developer.nordicsemi.com/nRF_Connect_SDK/doc/1.9.1/kconfig/CONFIG_CDC_ACM_IAD.html#help
    config.device_class = 0xEF;
    config.device_sub_class = 0x02;
    config.device_protocol = 0x01;
    config.composite_with_iads = true;

    config
}

/// This handles the low level USB management
#[embassy_executor::task]
pub async fn usb_task(mut usb: embassy_usb::UsbDevice<'static, AppDriver>) {
    usb.run().await;
}

#[embassy_executor::task]
async fn server_run(mut server: AppServer) {
    loop {
        // If the host disconnects, we'll return an error here.
        // If this happens, just wait until the host reconnects
        let _ = server.run().await;
    }
}

// ---

pub fn init_communication(usb_driver: Driver<'static, peripherals::USB>, impedance_setup: &'static ImpedanceSetupType, spawner: Spawner) {
    // Initialize communication peripherals
    let pbufs = PBUFS.take();
    let config = usb_config();

    let context = Context {
        unique_id: *uid::uid(),
        impedance_setup: impedance_setup,
    };

    let (device, tx_impl, rx_impl) = STORAGE.init(usb_driver, config, pbufs.tx_buf.as_mut_slice());
    let dispatcher = MyApp::new(context, spawner.into());
    let vkk = dispatcher.min_key_len();
    let server: AppServer = Server::new(
        tx_impl,
        rx_impl,
        pbufs.rx_buf.as_mut_slice(),
        dispatcher,
        vkk,
    );
    spawner.must_spawn(usb_task(device));
    spawner.must_spawn(server_run(server));
}

// Functions

pub fn ping_handler(_context: &mut Context, _header: VarHeader, rqst: u32) -> u32 {
    info!("ping");
    rqst
}

pub fn get_unique_id_handler(context: &mut Context, _header: VarHeader, _rqst: ()) -> [u8; 12] {
    info!("get_unique_id");
    context.unique_id
}

pub static LED_FREQUENCY_SIGNAL: Signal<CriticalSectionRawMutex, f32> = Signal::new();

pub async fn set_green_led_handler(_context: &mut Context, _header: VarHeader, rqst: f32) {
    info!("Set green led frequency to {:?} Hz", rqst);
    
    LED_FREQUENCY_SIGNAL.signal(rqst);
}

static STOP: Signal<CriticalSectionRawMutex, ()> = Signal::new();

#[embassy_executor::task]
pub async fn start_single_impedance_handler(context: SpawnCtx, header: VarHeader, rqst: SingleImpedanceStartRequest, sender: Sender<AppTx>) {
    info!("Start impedance measurement at {:?} Hz.", rqst.sinus_frequency);

    // Mark the impedance setup as running
    context.impedance_setup.lock().await.running_mode = RunningMode::SingleFrequency;

    // Init the sequencer
    let init_impedance_result = context.impedance_setup.lock().await.init_single_frequency_measurement(rqst.sinus_frequency, rqst.dft_number).await;
    // Trigger the sequencer
    context.impedance_setup.lock().await.start_measurement().await;

    if sender
        .reply::<StartSingleImpedanceEndpoint>(header.seq_no, &init_impedance_result)
        .await
        .is_err()
    {
        defmt::error!("Failed to reply, stopping accel");
        return;
    }

    let mut seq: u8 = 0;
    loop {
        let stop_fut = STOP.wait();
        let recv_fut = IMPEDANCE_CHANNEL_SINGLE.receive();

        match select(stop_fut, recv_fut).await {
            Either::First(_) => {
                info!("Stop signal received.");
                break;
            }
            Either::Second(msg) => {
                if sender
                    .publish::<SingleImpedanceOutputTopic>(seq.into(), &msg)
                    .await
                    .is_err()
                    {
                        defmt::error!("Topic send error!");
                        break;
                    }
                seq = seq.wrapping_add(1);
            }
        }
    }

    context.impedance_setup.lock().await.running_mode = RunningMode::None;
    info!("Impedance measurement stopped.");
    STOP.reset();
}

pub async fn stop_single_impedance_handler(context: &mut Context, _header: VarHeader, _rqst: ()) -> bool {
    info!("Stop impedance measurement");
    let was_busy = context.impedance_setup.lock().await.running_mode;
    if was_busy == RunningMode::SingleFrequency || was_busy == RunningMode::MultiFrequency(MeasurementPointSet::Eight) || was_busy == RunningMode::MultiFrequency(MeasurementPointSet::Eighteen) {
        STOP.signal(());
    }
    was_busy == RunningMode::SingleFrequency || was_busy == RunningMode::MultiFrequency(MeasurementPointSet::Eight) || was_busy == RunningMode::MultiFrequency(MeasurementPointSet::Eighteen)
}

#[embassy_executor::task]
pub async fn start_multi_impedance_handler(context: SpawnCtx, header: VarHeader, rqst: MultiImpedanceStartRequest, sender: Sender<AppTx>) {
    // Mark the impedance setup as running
    context.impedance_setup.lock().await.running_mode = RunningMode::MultiFrequency(rqst.points);
    
    // Init the sequencer
    let response = match rqst.points {
        MeasurementPointSet::Eight => {
            const SIZE: usize = 8;
            context
                .impedance_setup
                .lock()
                .await
                .init_multi_frequency_measurement::<SIZE>(rqst.points)
                .await
                .map(|periods| MultiImpedanceResult {
                    points: rqst.points,
                    periods_per_dft_8: periods,
                    periods_per_dft_18: Vec::new(),
                })
        }
        MeasurementPointSet::Eighteen => {
            const SIZE: usize = 18;
            context
                .impedance_setup
                .lock()
                .await
                .init_multi_frequency_measurement::<SIZE>(rqst.points)
                .await
                .map(|periods| MultiImpedanceResult {
                    points: rqst.points,
                    periods_per_dft_8: Vec::new(),
                    periods_per_dft_18: periods,
                })
        }
    };

    // Trigger the sequencer
    context.impedance_setup.lock().await.start_measurement().await;

    if sender
        .reply::<StartMultiImpedanceEndpoint>(header.seq_no, &response)
        .await
        .is_err()
    {
        defmt::error!("Failed to reply, stopping accel");
        return;
    }

    info!("Start multi impedance measurement.");

    let mut seq: u8 = 0;
    loop {
        let stop_fut = STOP.wait();
        let recv_fut = IMPEDANCE_CHANNEL_MULTI.receive();

        match select(stop_fut, recv_fut).await {
            Either::First(_) => {
                info!("Stop signal received.");
                break;
            }
            Either::Second(msg) => {
                if sender
                    .publish::<MultiImpedanceOutputTopic>(seq.into(), &msg)
                    .await
                    .is_err()
                    {
                        defmt::error!("Topic send error!");
                        break;
                    }
                seq = seq.wrapping_add(1);
            }
        }
    }

    context.impedance_setup.lock().await.running_mode = RunningMode::None;
    info!("Impedance measurement stopped.");
    STOP.reset();
}

pub async fn stop_multi_impedance_handler(context: &mut Context, _header: VarHeader, _rqst: ()) -> bool {
    info!("Stop impedance measurement");
    let was_busy = context.impedance_setup.lock().await.    running_mode;
    if was_busy == RunningMode::SingleFrequency || was_busy == RunningMode::MultiFrequency(MeasurementPointSet::Eight) || was_busy == RunningMode::MultiFrequency(MeasurementPointSet::Eighteen) {
        STOP.signal(());
    } 
    was_busy == RunningMode::SingleFrequency || was_busy == RunningMode::MultiFrequency(MeasurementPointSet::Eight) || was_busy == RunningMode::MultiFrequency(MeasurementPointSet::Eighteen)
}