autopilot/crates/gimbal_controller/src/internal/a40_protocol/frame.rs

//! Frame encoder / decoder for the ViewPro A40 vendor protocol.
//!
//! Wire format reminder (see module docs): `0x55 0xAA 0xDC` header,
//! length+counter byte, frame id, data, XOR checksum. We expose two
//! pure functions — [`encode_frame`] (Frame → bytes) and
//! [`decode_frame`] (bytes → Frame or [`FrameDecodeError`]).

use super::checksum::xor_checksum;

/// Vendor-fixed maximum packet size, including header (3) + length (1)
/// + frame id (1) + data + checksum (1). Anything larger is a protocol error.
pub const MAX_PACKET_LEN: usize = 63;

const HEADER_0: u8 = 0x55;
const HEADER_1: u8 = 0xAA;
const HEADER_2: u8 = 0xDC;
const HEADER_LEN: usize = 3;

/// Length-byte body-bits mask (bits 0..5).
const LENGTH_BODY_MASK: u8 = 0x3F;
/// Length-byte counter-bits shift (bits 6..7).
const LENGTH_COUNTER_SHIFT: u8 = 6;
/// Minimum body length (length byte + frame id + at least one data
/// byte + checksum = 4). Vendor SDK spec.
pub const MIN_BODY_LEN: u8 = 4;
/// Maximum body length (vendor SDK spec).
pub const MAX_BODY_LEN: u8 = 63;

/// Frame identifiers we use. Values are vendor-assigned and MUST NOT
/// be renumbered. See `AP_Mount_Viewpro.h::FrameId`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum FrameId {
    /// Handshake (sent to gimbal). Gimbal replies with T1_F1_B1_D1.
    Handshake = 0x00,
    /// Communication-config control (sent).
    U = 0x01,
    /// Communication-config status (received reply to U).
    V = 0x02,
    /// Heartbeat (received from gimbal).
    Heartbeat = 0x10,
    /// Target angles — yaw + pitch (sent).
    A1 = 0x1A,
    /// Camera controls, common (sent) — zoom in / zoom out / start
    /// record / stop record / take picture.
    C1 = 0x1C,
    /// Camera controls, less common (sent) — including absolute zoom
    /// (`CameraCommand2::SET_EO_ZOOM`).
    C2 = 0x2C,
    /// Tracking controls, common (sent).
    E1 = 0x1E,
    /// Tracking controls, less common (sent).
    E2 = 0x2E,
    /// Actual roll/pitch/yaw + recording/tracking status (received).
    T1F1B1D1 = 0x40,
    /// Vehicle attitude and position envelope (sent).
    Mahrs = 0xB1,
}

impl FrameId {
    pub fn from_u8(byte: u8) -> Option<Self> {
        match byte {
            0x00 => Some(Self::Handshake),
            0x01 => Some(Self::U),
            0x02 => Some(Self::V),
            0x10 => Some(Self::Heartbeat),
            0x1A => Some(Self::A1),
            0x1C => Some(Self::C1),
            0x2C => Some(Self::C2),
            0x1E => Some(Self::E1),
            0x2E => Some(Self::E2),
            0x40 => Some(Self::T1F1B1D1),
            0xB1 => Some(Self::Mahrs),
            _ => None,
        }
    }
}

/// Decoded frame. The frame-id field is canonicalised to the enum;
/// the data payload is the raw bytes that followed it in the wire
/// packet (excluding the length byte and the checksum).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Frame {
    pub frame_id: FrameId,
    pub data: Vec<u8>,
    /// Frame counter the sender stamped into bits 6..7 of the length
    /// byte. Echoed back so callers can correlate request/reply when
    /// the vendor protocol does not provide a separate sequence
    /// number. Range: 0..=3.
    pub frame_counter: u8,
}

#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum FrameDecodeError {
    #[error("buffer too small ({len} bytes; need at least 6)")]
    TooShort { len: usize },
    #[error("buffer too large ({len} bytes; max {max})")]
    TooLong { len: usize, max: usize },
    #[error("bad header bytes [{0:#04x} {1:#04x} {2:#04x}]; expected 55 AA DC")]
    BadHeader(u8, u8, u8),
    #[error("declared body length {declared} mismatches frame size {actual}")]
    BodyLengthMismatch { declared: u8, actual: usize },
    #[error("declared body length {0} out of range {min}..={max}", min = MIN_BODY_LEN, max = MAX_BODY_LEN)]
    BodyLengthOutOfRange(u8),
    #[error("unknown frame id {0:#04x}")]
    UnknownFrameId(u8),
    #[error("checksum mismatch: expected {expected:#04x}, got {actual:#04x}")]
    BadChecksum { expected: u8, actual: u8 },
}

/// Encode a frame for the wire.
///
/// `frame_counter` is masked to bits 0..1 and packed into bits 6..7
/// of the length byte (callers normally use a wrapping 0..=3 counter
/// owned by the transport).
///
/// Returns `None` if the resulting body length would exceed
/// [`MAX_BODY_LEN`] (the vendor's hard upper bound).
pub fn encode_frame(frame_id: FrameId, data: &[u8], frame_counter: u8) -> Option<Vec<u8>> {
    // Body length = length byte (1) + frame id (1) + data + checksum (1).
    let body_len_usize = 1 + 1 + data.len() + 1;
    if body_len_usize < MIN_BODY_LEN as usize || body_len_usize > MAX_BODY_LEN as usize {
        return None;
    }
    let body_len = body_len_usize as u8;

    let counter_bits = (frame_counter & 0b11) << LENGTH_COUNTER_SHIFT;
    let length_byte = (body_len & LENGTH_BODY_MASK) | counter_bits;

    let mut out = Vec::with_capacity(HEADER_LEN + body_len_usize);
    out.extend_from_slice(&[HEADER_0, HEADER_1, HEADER_2]);
    out.push(length_byte);
    out.push(frame_id as u8);
    out.extend_from_slice(data);

    // Checksum covers bytes 3..end-of-data. We have not pushed the
    // checksum yet, so the slice is exactly the bytes we want.
    let cs = xor_checksum(&out[HEADER_LEN..]);
    out.push(cs);
    Some(out)
}

/// Decode a frame from the wire. Returns `Err` for any header,
/// length, frame-id, or checksum violation — the caller (transport)
/// is responsible for counting these as `vendor_faults_total` and
/// dropping the frame.
pub fn decode_frame(buf: &[u8]) -> Result<Frame, FrameDecodeError> {
    if buf.len() < HEADER_LEN + MIN_BODY_LEN as usize {
        return Err(FrameDecodeError::TooShort { len: buf.len() });
    }
    if buf.len() > MAX_PACKET_LEN {
        return Err(FrameDecodeError::TooLong {
            len: buf.len(),
            max: MAX_PACKET_LEN,
        });
    }
    if buf[0] != HEADER_0 || buf[1] != HEADER_1 || buf[2] != HEADER_2 {
        return Err(FrameDecodeError::BadHeader(buf[0], buf[1], buf[2]));
    }
    let length_byte = buf[3];
    let body_len = length_byte & LENGTH_BODY_MASK;
    let frame_counter = length_byte >> LENGTH_COUNTER_SHIFT;
    if !(MIN_BODY_LEN..=MAX_BODY_LEN).contains(&body_len) {
        return Err(FrameDecodeError::BodyLengthOutOfRange(body_len));
    }
    // Body spans buf[3..3+body_len]. The total packet length is
    // header (3) + body_len.
    let expected_total = HEADER_LEN + body_len as usize;
    if buf.len() != expected_total {
        return Err(FrameDecodeError::BodyLengthMismatch {
            declared: body_len,
            actual: buf.len(),
        });
    }
    let frame_id_byte = buf[4];
    let frame_id =
        FrameId::from_u8(frame_id_byte).ok_or(FrameDecodeError::UnknownFrameId(frame_id_byte))?;
    let data_end = buf.len() - 1;
    let data = buf[5..data_end].to_vec();
    let actual_cs = buf[data_end];
    let expected_cs = xor_checksum(&buf[HEADER_LEN..data_end]);
    if expected_cs != actual_cs {
        return Err(FrameDecodeError::BadChecksum {
            expected: expected_cs,
            actual: actual_cs,
        });
    }
    Ok(Frame {
        frame_id,
        data,
        frame_counter,
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn round_trip_a1_yaw_command() {
        // Arrange — A1 (target angles) payload:
        // 1 byte ServoStatus + 2 bytes yaw BE + 2 bytes pitch BE + 4 bytes unused = 9 bytes data.
        // Yaw = 30° -> raw = 30/360 * 65536 ≈ 5461.
        let data = vec![0x0B, 0x15, 0x55, 0x00, 0x00, 0, 0, 0, 0];

        // Act
        let bytes = encode_frame(FrameId::A1, &data, 0).expect("encode");
        let decoded = decode_frame(&bytes).expect("decode");

        // Assert
        assert_eq!(decoded.frame_id, FrameId::A1);
        assert_eq!(decoded.data, data);
        assert_eq!(decoded.frame_counter, 0);
    }

    #[test]
    fn round_trip_c1_zoom_in() {
        // Arrange — C1 (camera command) payload: 2 BE bytes
        // (sensor_zoom_cmd_be). EO1 sensor (0x01) + CameraCommand::ZOOM_IN (0x09)
        // packs as one u16 BE; for this test we just check round-trip.
        let data = vec![0x01, 0x09];

        // Act
        let bytes = encode_frame(FrameId::C1, &data, 1).expect("encode");
        let decoded = decode_frame(&bytes).expect("decode");

        // Assert
        assert_eq!(decoded.frame_id, FrameId::C1);
        assert_eq!(decoded.data, data);
        assert_eq!(decoded.frame_counter, 1);
    }

    #[test]
    fn frame_counter_packs_and_unpacks() {
        // Arrange
        let data = vec![0xAA];

        // Act + Assert — counter wraps mod 4
        for counter in 0..4u8 {
            let bytes = encode_frame(FrameId::C1, &data, counter).unwrap();
            let decoded = decode_frame(&bytes).unwrap();
            assert_eq!(decoded.frame_counter, counter, "counter={counter}");
        }
        // High bits of the counter argument are masked off
        let bytes = encode_frame(FrameId::C1, &data, 0xFF).unwrap();
        let decoded = decode_frame(&bytes).unwrap();
        assert_eq!(decoded.frame_counter, 0b11);
    }

    #[test]
    fn corrupted_checksum_is_detected() {
        // Arrange
        let data = vec![0x01, 0x09];
        let mut bytes = encode_frame(FrameId::C1, &data, 0).unwrap();
        let last = bytes.len() - 1;
        bytes[last] ^= 0x01; // flip one bit

        // Act
        let err = decode_frame(&bytes).unwrap_err();

        // Assert
        assert!(matches!(err, FrameDecodeError::BadChecksum { .. }));
    }

    #[test]
    fn bad_header_rejected() {
        // Arrange — replace the magic header with 00 00 00
        let mut bytes = encode_frame(FrameId::C1, &[0x01, 0x09], 0).unwrap();
        bytes[0] = 0x00;
        bytes[1] = 0x00;
        bytes[2] = 0x00;

        // Act
        let err = decode_frame(&bytes).unwrap_err();

        // Assert
        assert!(matches!(err, FrameDecodeError::BadHeader(0, 0, 0)));
    }

    #[test]
    fn truncated_frame_rejected() {
        // Arrange
        let bytes = encode_frame(FrameId::C1, &[0x01, 0x09], 0).unwrap();
        let truncated = &bytes[..bytes.len() - 1];

        // Act
        let err = decode_frame(truncated).unwrap_err();

        // Assert
        assert!(matches!(err, FrameDecodeError::BodyLengthMismatch { .. }));
    }

    #[test]
    fn empty_data_falls_under_min_body_len() {
        // Arrange — empty data would mean body_len = 3 (length + frame id + checksum)
        // which is below MIN_BODY_LEN (4). encode_frame rejects.

        // Act
        let result = encode_frame(FrameId::C1, &[], 0);

        // Assert
        assert!(result.is_none());
    }

    #[test]
    fn oversize_data_rejected_by_encoder() {
        // Arrange — data large enough to overflow MAX_BODY_LEN
        let data = vec![0; MAX_BODY_LEN as usize];

        // Act
        let result = encode_frame(FrameId::C1, &data, 0);

        // Assert
        assert!(result.is_none());
    }

    #[test]
    fn unknown_frame_id_rejected() {
        // Arrange — manually craft a frame with frame_id = 0x99
        let data = vec![0x01, 0x09];
        let bytes_ok = encode_frame(FrameId::C1, &data, 0).unwrap();
        let mut bytes = bytes_ok.clone();
        bytes[4] = 0x99; // overwrite frame id
                         // Recompute checksum so the decoder gets to the frame-id check
        let cs_idx = bytes.len() - 1;
        bytes[cs_idx] = xor_checksum(&bytes[3..cs_idx]);

        // Act
        let err = decode_frame(&bytes).unwrap_err();

        // Assert
        assert!(matches!(err, FrameDecodeError::UnknownFrameId(0x99)));
    }
}