protoflight 0.1.2

#![cfg(feature = "blackbox")]

use blackbox_logger::{Blackbox, BlackboxConfig, BlackboxMainData, BlackboxSlowData, LoggerState, SliceEncoder};
use embassy_sync::{blocking_mutex::raw::CriticalSectionRawMutex, channel::Channel};

use crate::{
    sensors::{GyroPidMessage, SetpointMessage},
    tasks::gyro_pid_task::{GyroPidReceiver, SetpointReceiver},
};

#[cfg(feature = "debug")]
use crate::tasks::DebugMode;

#[cfg(feature = "gps")]
use {
    crate::{
        gps::{GpsMessage, GpsSolutionData},
        tasks::gps_task::GpsSubscriber,
    },
    blackbox_logger::{BlackboxGpsData, BlackboxGpsPosition},
};

pub struct BlackboxContext {
    pub gyro_pid_receiver: GyroPidReceiver,
    pub setpoint_receiver: SetpointReceiver,
    pub setpoint_message: SetpointMessage,
    #[cfg(feature = "gps")]
    pub gps_subscriber: GpsSubscriber,
    pub blackbox: Blackbox,
    pub buffer: [u8; BlackboxContext::BUFFER_CAPACITY],
    pub overflow_counter: u32,
}

impl BlackboxContext {
    const BUFFER_CAPACITY: usize = 1024;

    #[rustfmt::skip]
    pub fn new(
        gyro_pid_receiver: GyroPidReceiver,
        setpoint_receiver: SetpointReceiver,
        setpoint_message: SetpointMessage,
        blackbox_config: BlackboxConfig,
        #[cfg(feature = "gps")] gps_subscriber: GpsSubscriber,
    ) -> Self {
        //let mut blackbox_config = blackbox_config;
        //blackbox_config.huffman_compress = true;

        // NRVO (Named Return Value Optimization) ensures blackbox is created in place and not copied.
        let mut blackbox = Blackbox::new(blackbox_config);
        blackbox.init();

        Self {
            gyro_pid_receiver,
            setpoint_receiver,
            setpoint_message,
            #[cfg(feature = "gps")] gps_subscriber,
            blackbox,
            buffer: [0u8; Self::BUFFER_CAPACITY],
            overflow_counter: 0,
        }
    }
}

/// A fixed-size message container used to pass blackbox chunks between tasks.
/// Set the internal buffer capacity to a size larger than your maximum possible `len` frame.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct BlackboxWriteBlock {
    pub data: [u8; Self::CAPACITY], // Adjust size to match your largest expected serialized packet length
    pub len: usize,
}

impl Default for BlackboxWriteBlock {
    fn default() -> Self {
        Self::new(0)
    }
}

impl BlackboxWriteBlock {
    pub const CAPACITY: usize = 64;

    #[inline]
    pub const fn new(len: usize) -> Self {
        Self { data: [0u8; Self::CAPACITY], len }
    }
    #[inline]
    pub fn from_chunk(slice: &[u8]) -> Self {
        // Enforce the size boundary strictly using compile-time constants
        let copy_len = core::cmp::min(slice.len(), Self::CAPACITY);
        let mut block = Self { data: [0u8; Self::CAPACITY], len: copy_len };
        block.data[..copy_len].copy_from_slice(&slice[..copy_len]);
        block
    }
}

const BLACKBOX_WRITE_QUEUE_COUNT: usize = 256;
pub static BLACKBOX_WRITE_QUEUE: Channel<CriticalSectionRawMutex, BlackboxWriteBlock, BLACKBOX_WRITE_QUEUE_COUNT> =
    Channel::new();

fn send_data_to_blackbox_writer_task(data: &[u8], overflow_counter: &mut u32) {
    _ = overflow_counter;
    // Loop through the slice in chunks matching BlackboxWriteBlock capacity
    for chunk in data.chunks(BlackboxWriteBlock::CAPACITY) {
        let block = BlackboxWriteBlock::from_chunk(chunk);
        // Non-blocking try_send ensures high-speed loop deadlines are protected
        if let Err(_overflow) = BLACKBOX_WRITE_QUEUE.try_send(block) {
            *overflow_counter = overflow_counter.wrapping_add(1);
            log::error!("BLACKBOX: FIFO queue full! Dropped a log chunk.");
        }
    }
}

/// Blackbox task.
#[embassy_executor::task]
pub async fn blackbox_task(ctx: &'static mut BlackboxContext) {
    log::info!("    BLACKBOX: task started");
    let mut loop_count: u32 = 0;

    // Write the Blackbox log file header by using blackbox.update to step through the blackbox state machine
    // until the state is LoggerState::HeaderWritten.
    ctx.blackbox.set_state(LoggerState::WriteFileHeader);
    while ctx.blackbox.state() != LoggerState::HeaderWritten {
        let time_us = 0;
        let len = ctx.blackbox.update(&mut SliceEncoder::new(&mut ctx.buffer), time_us);
        send_data_to_blackbox_writer_task(&ctx.buffer[..len], &mut ctx.overflow_counter);
        //log::info!("BLACKBOX:  hdr {loop_count},{len}");
        loop_count = loop_count.wrapping_add(1);
    }
    log::info!("    BLACKBOX: header written {loop_count}");

    loop_count = 0;
    loop {
        // blocking
        let gyro_pid_msg = ctx.gyro_pid_receiver.changed().await;
        let time_us = gyro_pid_msg.time_us;
        // non-blocking
        if let Some(setpoint_msg) = ctx.setpoint_receiver.try_get() {
            // if we have a new setpoint message then update ctx.setpoint_message so that the most up to date setpoint_message is used.
            ctx.setpoint_message = setpoint_msg;
            let mut slow_data = slow_data_from(ctx.setpoint_message);
            slow_data.set_blackbox_active(true);
            ctx.blackbox.set_slow_data(slow_data);
        }
        // set_main_data always uses the most up to date setpoint message.
        ctx.blackbox.set_main_data(main_data_from(gyro_pid_msg, ctx.setpoint_message));

        #[cfg(feature = "gps")]
        if let Some(wait_result) = ctx.gps_subscriber.try_next_message()
            && let embassy_sync::pubsub::WaitResult::Message(event) = wait_result
            && let GpsMessage::GpsSolution(gps_solution_data) = event
        {
            let gps_data = gps_data_from(gps_solution_data);
            //let satellite_count = gps_data.satellite_count;
            //log::info!("    BLACKBOX: sat count {satellite_count}");
            ctx.blackbox.set_gps_data(gps_data);
        }

        let len = ctx.blackbox.update(&mut SliceEncoder::new(&mut ctx.buffer), time_us);
        /*#[cfg(feature = "std")]
        if loop_count == 512 {
            // write End of log
            let len = ctx.blackbox.logger.log_e_frame(&mut SliceEncoder::new(&mut ctx.buffer), BlackboxEvent::LogEnd);
            send_data_to_blackbox_writer_task(&ctx.buffer[..len], &mut ctx.overflow_counter);
            log::info!("**** BLACKBOX: END OF LOG");
        }*/
        send_data_to_blackbox_writer_task(&ctx.buffer[..len], &mut ctx.overflow_counter);
        if loop_count.is_multiple_of(10) {
            log::info!("      BLACKBOX: loop {loop_count},{len}");
        }
        loop_count = loop_count.wrapping_add(1); // use wrapping_add to handle when time rolls over at max u32.
    }

    /*loop {
        // Wait until the sender updates the value.
        // The primary client of the ahrs_receiver is the motor_mixer
        let gyro_pid_msg = ctx.gyro_pid_receiver.changed().await;
        ctx.blackbox.load_telemetry(time_us, gyro_pid_msg, setpoint_msg);

        log::info!("BLACKBOX: Received time_us {}", gyro_pid_msg.time_us);
        let len = 0;
        log::info!(
            "BLACKBOX: Encoded frame in {} bytes. (x: {}, y: {}, z: {}),(x: {}, y: {}, z: {})",
            len,
            gyro_pid_msg.gyro_rps.x,
            gyro_pid_msg.gyro_rps.y,
            gyro_pid_msg.gyro_rps.z,
            gyro_pid_msg.gyro_rps_unfiltered.x,
            gyro_pid_msg.gyro_rps_unfiltered.y,
            gyro_pid_msg.gyro_rps_unfiltered.z
        );
        // Process logging.
        /*let buf = [b'a', b'b', b'c', b'd', b'e', b'f'];
        let len = 6;
        ctx.sd_card.write_all(&buf[..len]).await;*/
        // 3. Increment fake time (e.g., 1000us per sample for 1kHz)
        time_us = time_us.wrapping_add(1000); // use wrapping_add to handle when time rolls over at max u32.
    }*/
}

#[inline]
pub fn main_data_from(gyro_pid_msg: GyroPidMessage, _setpoint_msg: SetpointMessage) -> BlackboxMainData {
    const TO_I16: f32 = 32_757.0;

    #[cfg(feature = "debug")]
    crate::tasks::GLOBAL_DEBUG.set(DebugMode::BlackboxOutput, 0, 0);

    debug_set_mode!(DebugMode::BlackboxOutput);
    debug_set!(DebugMode::BlackboxOutput, 0, 0);

    let motor_commands = gyro_pid_msg.motor_commands * 2.0;
    BlackboxMainData {
        time_us: gyro_pid_msg.time_us,
        baro_altitude: 0,
        range_raw: 0,
        amperage: 0,
        battery_voltage: 0,
        rssi: 0,
        // todo, add scaling to below
        #[allow(clippy::cast_possible_truncation)]
        pid_p: gyro_pid_msg.pid_errors_p.map(|x| x as i32),
        #[allow(clippy::cast_possible_truncation)]
        pid_i: gyro_pid_msg.pid_errors_i.map(|x| x as i32),
        #[allow(clippy::cast_possible_truncation)]
        pid_d: [gyro_pid_msg.pid_errors_d[0] as i32, gyro_pid_msg.pid_errors_d[1] as i32, 0],
        pid_s: [0i32; BlackboxMainData::RPY_AXIS_COUNT],
        pid_k: [0i32; BlackboxMainData::RPY_AXIS_COUNT],

        rc_commands: [1500, 1500, 1500, 1100],

        // TODO: need to scale these
        #[allow(clippy::cast_possible_truncation)]
        setpoints: [motor_commands.x as i16, motor_commands.y as i16, motor_commands.z as i16, motor_commands.t as i16],
        gyro: (gyro_pid_msg.gyro_rps.to_degrees()).into(),
        gyro_unfiltered: (gyro_pid_msg.gyro_rps_unfiltered.to_degrees()).into(),
        acc: (gyro_pid_msg.acc * 4096.0).into(),
        #[cfg(feature = "magnetometer")]
        mag: [0i16; BlackboxMainData::XYZ_AXIS_COUNT],

        #[allow(clippy::cast_possible_truncation)]
        orientation: if gyro_pid_msg.orientation.w > 0.0 {
            [
                (gyro_pid_msg.orientation.x * TO_I16) as i16,
                (gyro_pid_msg.orientation.y * TO_I16) as i16,
                (gyro_pid_msg.orientation.z * TO_I16) as i16,
            ]
        } else {
            [
                (-gyro_pid_msg.orientation.x * TO_I16) as i16,
                (-gyro_pid_msg.orientation.y * TO_I16) as i16,
                (-gyro_pid_msg.orientation.z * TO_I16) as i16,
            ]
        },
        #[cfg(feature = "eight_motors")]
        motor: [1100, 1100, 1100, 1100, 1100, 1100, 1100, 1100],
        #[cfg(not(feature = "eight_motors"))]
        motor: [1100, 1100, 1100, 1100],
        #[cfg(feature = "dshot_telemetry")]
        erpm: [0u16; BlackboxMainData::MAX_SUPPORTED_MOTOR_COUNT],

        #[cfg(not(feature = "debug"))]
        debug: [0i16; BlackboxMainData::DEBUG_COUNT],
        #[cfg(feature = "debug")]
        debug: crate::tasks::GLOBAL_DEBUG.values(),

        #[cfg(feature = "servos")]
        servos: [0i16; MainData::MAX_SUPPORTED_SERVO_COUNT],
    }
}

#[inline]
pub fn slow_data_from(setpoint: SetpointMessage) -> BlackboxSlowData {
    BlackboxSlowData {
        flight_mode_flags: setpoint.rc_modes.bits_0_31(),
        gps_state_flags: setpoint.gps_state_flags,
        failsafe_phase: setpoint.failsafe_phase,
        rx_signal_received: setpoint.rx_signal_received,
        rx_flight_channel_is_valid: setpoint.rx_flight_channel_is_valid,
    }
}

#[cfg(feature = "gps")]
#[inline]
pub fn gps_data_from(gps: GpsSolutionData) -> BlackboxGpsData {
    BlackboxGpsData {
        time_of_week_ms: gps.time,
        interval_ms: 0,
        position: BlackboxGpsPosition {
            longitude_degrees_x1e7: gps.llh.longitude_degrees_x1e7,
            latitude_degrees_x1e7: gps.llh.latitude_degrees_x1e7,
            altitude_cm: gps.llh.altitude_cm,
        },
        velocity_north_cmps: gps.velocity_ned_cmps.north,
        velocity_east_cmps: gps.velocity_ned_cmps.east,
        velocity_down_cmps: gps.velocity_ned_cmps.down,
        speed3d_cmps: gps.speed3d_cmps.cast_signed(),
        ground_speed_cmps: gps.ground_speed_cmps.cast_signed(),
        ground_course_degrees_x10: gps.ground_course_degrees_x10.cast_signed(),
        satellite_count: gps.satellite_count,
    }
}