xblink/src/ntag5/sram.rs

//! SRAM mailbox command/response protocol for NFC<->MCU communication.
//!
//! The phone writes a command packet to NTAG5Link SRAM via NFC.
//! The MCU reads it over I2C, processes the command, and writes
//! a response packet back to SRAM for the phone to read.

use crate::ntag5::{self, Ntag5Link, SRAM_SIZE};
use crate::pattern;
use embedded_hal::i2c::I2c;

// ---------------------------------------------------------------------------
// Constants
// ---------------------------------------------------------------------------

/// Firmware version reported in GET_STATUS response.
pub const FIRMWARE_VERSION: u8 = 0x01;

// Command IDs (phone -> MCU)
pub const CMD_WRITE_PATTERN: u8 = 0x01;
pub const CMD_GET_STATUS: u8 = 0x02;
pub const CMD_SET_ACTIVE: u8 = 0x03;
pub const CMD_SYNC_START: u8 = 0x05;
pub const CMD_SYNC_END: u8 = 0x06;
pub const CMD_READ_LIBRARY: u8 = 0x07;
pub const CMD_READ_NEXT: u8 = 0x08;

// Response marker
pub const RESPONSE_MARKER: u8 = 0xFF;

// Status codes
pub const STATUS_OK: u8 = 0x00;
pub const STATUS_BAD_CMD: u8 = 0x02;
pub const STATUS_EEPROM_FAIL: u8 = 0x03;
pub const STATUS_INVALID_INDEX: u8 = 0x04;

/// Command header size (cmd + seq + payload_len_u16).
pub const CMD_HEADER_SIZE: usize = 4;

/// Response header size (marker + seq + status + payload_len).
pub const RSP_HEADER_SIZE: usize = 4;

// ---------------------------------------------------------------------------
// Mailbox state
// ---------------------------------------------------------------------------

/// Persistent state for the mailbox protocol across commands.
pub struct MailboxState {
    /// True while a SYNC_START..SYNC_END sequence is in progress.
    pub syncing: bool,
    /// Pattern count supplied by SYNC_START.
    pub sync_count: u8,
    /// Current (mA setting) supplied by SYNC_START.
    pub sync_current: u8,
    /// LED mode supplied by SYNC_START.
    pub sync_mode: u8,
    /// Next pattern index for READ_NEXT.
    pub read_index: u8,
    /// Total patterns available for READ_NEXT iteration.
    pub read_count: u8,
}

impl MailboxState {
    pub fn new() -> Self {
        Self {
            syncing: false,
            sync_count: 0,
            sync_current: 0,
            sync_mode: 0,
            read_index: 0,
            read_count: 0,
        }
    }
}

// ---------------------------------------------------------------------------
// Protocol parse / build
// ---------------------------------------------------------------------------

/// Parse the 4-byte command header. Returns (cmd, seq, payload_len) or None
/// if the buffer is too short or the command ID is out of range.
fn parse_header(buf: &[u8]) -> Option<(u8, u8, u16)> {
    if buf.len() < CMD_HEADER_SIZE {
        return None;
    }
    let cmd = buf[0];
    if cmd == 0 || cmd >= 0x80 {
        return None;
    }
    let seq = buf[1];
    let payload_len = (buf[2] as u16) | ((buf[3] as u16) << 8); // LE
    Some((cmd, seq, payload_len))
}

/// Build a response packet into `buf`. Returns the total number of bytes
/// written (padded to a multiple of 4 for SRAM block alignment).
fn build_response(buf: &mut [u8], seq: u8, status: u8, payload: &[u8]) -> usize {
    let plen = payload.len();
    buf[0] = RESPONSE_MARKER;
    buf[1] = seq;
    buf[2] = status;
    buf[3] = plen as u8;
    buf[RSP_HEADER_SIZE..RSP_HEADER_SIZE + plen].copy_from_slice(payload);
    // Pad total length to multiple of 4
    let raw_len = RSP_HEADER_SIZE + plen;
    let padded = (raw_len + 3) & !3;
    for b in buf[raw_len..padded].iter_mut() {
        *b = 0;
    }
    padded
}

// ---------------------------------------------------------------------------
// EEPROM helpers
// ---------------------------------------------------------------------------

/// Block address for pattern N in the EEPROM library.
fn pattern_block(index: u8) -> u16 {
    pattern::LIBRARY_BASE_BLOCK + 4 + (index as u16) * 28
}

/// Read the XBLK library header from EEPROM.
/// Returns None if magic/version don't match.
fn read_library_header<I2C, E>(
    ntag: &mut Ntag5Link<I2C>,
) -> Result<Option<pattern::LibraryHeader>, ntag5::Error<E>>
where
    I2C: I2c<Error = E>,
{
    let mut hdr_buf = [0u8; pattern::HEADER_SIZE];
    ntag.read_memory(pattern::LIBRARY_BASE_BLOCK, &mut hdr_buf)?;
    Ok(pattern::parse_header(&hdr_buf))
}

/// Write the 16-byte header to EEPROM (4 blocks starting at LIBRARY_BASE_BLOCK).
fn write_header_to_eeprom<I2C, E>(
    ntag: &mut Ntag5Link<I2C>,
    hdr: &[u8; pattern::HEADER_SIZE],
    delay: &mut impl embedded_hal::delay::DelayNs,
) -> Result<(), ntag5::Error<E>>
where
    I2C: I2c<Error = E>,
{
    for i in 0..4u16 {
        let mut chunk = [0u8; 4];
        let off = (i as usize) * 4;
        chunk.copy_from_slice(&hdr[off..off + 4]);
        ntag.write_verify_block(pattern::LIBRARY_BASE_BLOCK + i, &chunk, delay)?;
    }
    Ok(())
}

/// Read all pattern data from EEPROM for the given count (for CRC computation).
/// Returns the number of bytes read into `all_data`.
fn read_all_pattern_data<I2C, E>(
    ntag: &mut Ntag5Link<I2C>,
    count: u8,
    all_data: &mut [u8],
) -> Result<usize, ntag5::Error<E>>
where
    I2C: I2c<Error = E>,
{
    let total = count as usize * pattern::PATTERN_ENTRY_SIZE;
    // Read in chunks — read_memory can handle arbitrary lengths
    // but we read per-pattern for clarity
    for i in 0..count as usize {
        let block = pattern_block(i as u8);
        let off = i * pattern::PATTERN_ENTRY_SIZE;
        ntag.read_memory(block, &mut all_data[off..off + pattern::PATTERN_ENTRY_SIZE])?;
    }
    Ok(total)
}

/// Recalculate the header CRC based on current EEPROM pattern data,
/// then write the updated header back.
fn recalculate_header_crc<I2C, E>(
    ntag: &mut Ntag5Link<I2C>,
    count: u8,
    active: u8,
    current: u8,
    led_mode: u8,
    delay: &mut impl embedded_hal::delay::DelayNs,
) -> Result<(), ntag5::Error<E>>
where
    I2C: I2c<Error = E>,
{
    // Read all pattern data for CRC
    let mut all_data = [0u8; pattern::MAX_PATTERNS * pattern::PATTERN_ENTRY_SIZE];
    let data_len = read_all_pattern_data(ntag, count, &mut all_data)?;

    let mut hdr = [0u8; pattern::HEADER_SIZE];
    pattern::serialize_header(count, active, current, led_mode, &all_data[..data_len], &mut hdr);
    write_header_to_eeprom(ntag, &hdr, delay)
}

/// After a standalone WRITE_PATTERN (not during sync), re-read the header
/// and recalculate its CRC to account for the new pattern data.
fn update_header_after_write<I2C, E>(
    ntag: &mut Ntag5Link<I2C>,
    delay: &mut impl embedded_hal::delay::DelayNs,
) -> Result<(), ntag5::Error<E>>
where
    I2C: I2c<Error = E>,
{
    let hdr = match read_library_header(ntag)? {
        Some(h) => h,
        None => return Ok(()), // no valid header, nothing to update
    };
    recalculate_header_crc(
        ntag,
        hdr.pattern_count,
        hdr.active_index,
        hdr.current,
        hdr.led_mode,
        delay,
    )
}

// ---------------------------------------------------------------------------
// Command handlers
// ---------------------------------------------------------------------------

/// GET_STATUS (0x02): Return firmware version + library summary.
fn handle_get_status<I2C, E>(
    ntag: &mut Ntag5Link<I2C>,
    seq: u8,
    rsp_buf: &mut [u8],
) -> Result<usize, ntag5::Error<E>>
where
    I2C: I2c<Error = E>,
{
    let mut payload = [0u8; 5];
    payload[0] = FIRMWARE_VERSION;
    match read_library_header(ntag)? {
        Some(h) => {
            payload[1] = h.pattern_count;
            payload[2] = h.active_index;
            payload[3] = h.current;
            payload[4] = h.led_mode;
        }
        None => {
            // No valid header — return zeros
        }
    }
    Ok(build_response(rsp_buf, seq, STATUS_OK, &payload))
}

/// WRITE_PATTERN (0x01): Write a 112-byte pattern entry to EEPROM.
fn handle_write_pattern<I2C, E>(
    ntag: &mut Ntag5Link<I2C>,
    state: &MailboxState,
    seq: u8,
    payload: &[u8],
    rsp_buf: &mut [u8],
    delay: &mut impl embedded_hal::delay::DelayNs,
) -> Result<usize, ntag5::Error<E>>
where
    I2C: I2c<Error = E>,
{
    if payload.len() < 1 + pattern::PATTERN_ENTRY_SIZE {
        return Ok(build_response(rsp_buf, seq, STATUS_INVALID_INDEX, &[]));
    }
    let index = payload[0];
    if index as usize >= pattern::MAX_PATTERNS {
        return Ok(build_response(rsp_buf, seq, STATUS_INVALID_INDEX, &[]));
    }

    // Write 112 bytes = 28 blocks
    let base_block = pattern_block(index);
    let pattern_data = &payload[1..1 + pattern::PATTERN_ENTRY_SIZE];
    for i in 0..28u16 {
        let off = (i as usize) * 4;
        let mut chunk = [0u8; 4];
        chunk.copy_from_slice(&pattern_data[off..off + 4]);
        match ntag.write_verify_block(base_block + i, &chunk, delay) {
            Ok(()) => {}
            Err(ntag5::Error::VerifyFailed) => {
                return Ok(build_response(rsp_buf, seq, STATUS_EEPROM_FAIL, &[]));
            }
            Err(e) => return Err(e),
        }
    }

    // If not syncing, update header CRC to reflect changed pattern data
    if !state.syncing {
        match update_header_after_write(ntag, delay) {
            Ok(()) => {}
            Err(ntag5::Error::VerifyFailed) => {
                return Ok(build_response(rsp_buf, seq, STATUS_EEPROM_FAIL, &[]));
            }
            Err(e) => return Err(e),
        }
    }

    Ok(build_response(rsp_buf, seq, STATUS_OK, &[]))
}

/// SET_ACTIVE (0x03): Change the active pattern index in the EEPROM header.
fn handle_set_active<I2C, E>(
    ntag: &mut Ntag5Link<I2C>,
    seq: u8,
    payload: &[u8],
    rsp_buf: &mut [u8],
    delay: &mut impl embedded_hal::delay::DelayNs,
) -> Result<usize, ntag5::Error<E>>
where
    I2C: I2c<Error = E>,
{
    if payload.is_empty() {
        return Ok(build_response(rsp_buf, seq, STATUS_INVALID_INDEX, &[]));
    }
    let index = payload[0];

    let hdr = match read_library_header(ntag)? {
        Some(h) => h,
        None => return Ok(build_response(rsp_buf, seq, STATUS_EEPROM_FAIL, &[])),
    };

    if index >= hdr.pattern_count {
        return Ok(build_response(rsp_buf, seq, STATUS_INVALID_INDEX, &[]));
    }

    // Recalculate header with new active index
    match recalculate_header_crc(ntag, hdr.pattern_count, index, hdr.current, hdr.led_mode, delay)
    {
        Ok(()) => {}
        Err(ntag5::Error::VerifyFailed) => {
            return Ok(build_response(rsp_buf, seq, STATUS_EEPROM_FAIL, &[]));
        }
        Err(e) => return Err(e),
    }

    Ok(build_response(rsp_buf, seq, STATUS_OK, &[]))
}

/// SYNC_START (0x05): Begin a bulk pattern upload session.
fn handle_sync_start(
    state: &mut MailboxState,
    seq: u8,
    payload: &[u8],
    rsp_buf: &mut [u8],
) -> usize {
    if payload.len() < 3 {
        return build_response(rsp_buf, seq, STATUS_BAD_CMD, &[]);
    }
    let count = payload[0];
    if count == 0 || count as usize > pattern::MAX_PATTERNS {
        return build_response(rsp_buf, seq, STATUS_INVALID_INDEX, &[]);
    }
    state.syncing = true;
    state.sync_count = count;
    state.sync_current = payload[1];
    state.sync_mode = payload[2];
    build_response(rsp_buf, seq, STATUS_OK, &[])
}

/// SYNC_END (0x06): Finalize bulk upload — write the XBLK header with CRC.
fn handle_sync_end<I2C, E>(
    ntag: &mut Ntag5Link<I2C>,
    state: &mut MailboxState,
    seq: u8,
    rsp_buf: &mut [u8],
    delay: &mut impl embedded_hal::delay::DelayNs,
) -> Result<usize, ntag5::Error<E>>
where
    I2C: I2c<Error = E>,
{
    if !state.syncing {
        return Ok(build_response(rsp_buf, seq, STATUS_BAD_CMD, &[]));
    }

    // Active index defaults to 0
    match recalculate_header_crc(
        ntag,
        state.sync_count,
        0, // active_index = 0
        state.sync_current,
        state.sync_mode,
        delay,
    ) {
        Ok(()) => {}
        Err(ntag5::Error::VerifyFailed) => {
            state.syncing = false;
            return Ok(build_response(rsp_buf, seq, STATUS_EEPROM_FAIL, &[]));
        }
        Err(e) => {
            state.syncing = false;
            return Err(e);
        }
    }

    state.syncing = false;
    Ok(build_response(rsp_buf, seq, STATUS_OK, &[]))
}

/// READ_LIBRARY (0x07): Return library summary and prepare for READ_NEXT.
fn handle_read_library<I2C, E>(
    ntag: &mut Ntag5Link<I2C>,
    state: &mut MailboxState,
    seq: u8,
    rsp_buf: &mut [u8],
) -> Result<usize, ntag5::Error<E>>
where
    I2C: I2c<Error = E>,
{
    let mut payload = [0u8; 4];
    match read_library_header(ntag)? {
        Some(h) => {
            payload[0] = h.pattern_count;
            payload[1] = h.active_index;
            payload[2] = h.current;
            payload[3] = h.led_mode;
            state.read_index = 0;
            state.read_count = h.pattern_count;
        }
        None => {
            state.read_index = 0;
            state.read_count = 0;
        }
    }
    Ok(build_response(rsp_buf, seq, STATUS_OK, &payload))
}

/// READ_NEXT (0x08): Return the next pattern entry (112 bytes).
fn handle_read_next<I2C, E>(
    ntag: &mut Ntag5Link<I2C>,
    state: &mut MailboxState,
    seq: u8,
    rsp_buf: &mut [u8],
) -> Result<usize, ntag5::Error<E>>
where
    I2C: I2c<Error = E>,
{
    if state.read_index >= state.read_count {
        return Ok(build_response(rsp_buf, seq, STATUS_INVALID_INDEX, &[]));
    }

    let mut entry = [0u8; pattern::PATTERN_ENTRY_SIZE];
    let block = pattern_block(state.read_index);
    ntag.read_memory(block, &mut entry)?;
    state.read_index += 1;
    Ok(build_response(rsp_buf, seq, STATUS_OK, &entry))
}

// ---------------------------------------------------------------------------
// Main entry point
// ---------------------------------------------------------------------------

/// Read SRAM, dispatch the command, and write the response back.
///
/// Returns `Ok(true)` if a command was processed, `Ok(false)` if no valid
/// command was found in SRAM (e.g., empty buffer, bad header, bad CRC).
///
/// I2C errors propagate as `Err`. EEPROM verify failures are reported via
/// a STATUS_EEPROM_FAIL response (not as Err).
pub fn process_command<I2C, E>(
    ntag: &mut Ntag5Link<I2C>,
    state: &mut MailboxState,
    delay: &mut impl embedded_hal::delay::DelayNs,
) -> Result<bool, ntag5::Error<E>>
where
    I2C: I2c<Error = E>,
{
    // Read full SRAM
    let mut sram = [0u8; SRAM_SIZE];
    ntag.read_sram(&mut sram)?;

    // Parse header
    let (cmd, seq, payload_len) = match parse_header(&sram) {
        Some(h) => h,
        None => return Ok(false),
    };

    // Bounds check
    let total = CMD_HEADER_SIZE + payload_len as usize;
    if total > SRAM_SIZE {
        return Ok(false);
    }

    let payload = &sram[CMD_HEADER_SIZE..total];

    // Dispatch
    let mut rsp = [0u8; SRAM_SIZE];
    let rsp_len = match cmd {
        CMD_GET_STATUS => handle_get_status(ntag, seq, &mut rsp)?,
        CMD_WRITE_PATTERN => {
            handle_write_pattern(ntag, state, seq, payload, &mut rsp, delay)?
        }
        CMD_SET_ACTIVE => handle_set_active(ntag, seq, payload, &mut rsp, delay)?,
        CMD_SYNC_START => handle_sync_start(state, seq, payload, &mut rsp),
        CMD_SYNC_END => handle_sync_end(ntag, state, seq, &mut rsp, delay)?,
        CMD_READ_LIBRARY => handle_read_library(ntag, state, seq, &mut rsp)?,
        CMD_READ_NEXT => handle_read_next(ntag, state, seq, &mut rsp)?,
        _ => build_response(&mut rsp, seq, STATUS_BAD_CMD, &[]),
    };

    ntag.write_sram_blocks(0, &rsp[..rsp_len])?;
    Ok(true)
}