uftwo 0.3.0

use core::fmt;
use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout};

/// Block size in bytes (UF2 specification).
pub const BLOCK_SIZE: usize = 512;

/// Maximum payload size in a UF2 block.
pub const MAX_PAYLOAD_SIZE: usize = 476;

/// Size of checksum placed within payload
const CHECKSUM_SIZE: usize = 24;

/// Maximum payload size when checksum is included (476 - 24 = 452).
pub const MAX_PAYLOAD_SIZE_WITH_CHECKSUM: usize =
    MAX_PAYLOAD_SIZE - CHECKSUM_SIZE;

/// Padding byte used in UF2 blocks.
pub const PADDING_BYTE: u8 = 0xFF;

/// Align to 4 byte boundary.
pub const ALIGN: usize = 4;

/// Magic numbers.
pub const MAGIC_NUMBER: [u32; 3] = [0x0A324655, 0x9E5D5157, 0x0AB16F30];

/// Block error kind.
#[derive(Debug, PartialEq, Eq)]
#[cfg_attr(feature = "defmt", derive(defmt::Format))]
pub enum BlockError {
    /// There was an issue with the input buffer size or alignment.
    InputBuffer,
    /// One or more of the magic numbers were incorrect.
    MagicNumber,
    /// Payload size too large.
    PayloadSize,
    /// Block number invalid, e.g. exceeds total blocks count.
    BlockNumberInvalid,
}

impl fmt::Display for BlockError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::InputBuffer => write!(f, "Input buffer"),
            Self::MagicNumber => write!(f, "Magic number incorrect"),
            Self::PayloadSize => write!(f, "Payload size too large"),
            Self::BlockNumberInvalid => write!(f, "Block number invalid"),
        }
    }
}

impl core::error::Error for BlockError {}

/// Block structure.
///
/// Length is fixed at 512 bytes with a variable size data section up to 476 bytes.
#[derive(Debug, Copy, Clone, Immutable, KnownLayout, FromBytes, IntoBytes)]
#[repr(C)]
#[cfg_attr(feature = "defmt", derive(defmt::Format))]
pub struct Block {
    /// First magic number.
    pub magic_start_0: u32,
    /// Second magic number.
    pub magic_start_1: u32,
    /// Flags.
    pub flags: Flags,
    /// Address in flash where the data should be written.
    pub target_addr: u32,
    /// Number of bytes used in data.
    pub data_len: u32,
    /// Sequential block number, starting at 0.
    pub block: u32,
    /// Total number of blocks.
    pub total_blocks: u32,
    /// File size or board family ID or zero.
    pub board_family_id_or_file_size: u32,
    /// Payload data, padded with zeros.
    ///
    /// When the MD5 checksum flag is set, the last 24 bytes hold the checksum
    /// as well as address start and length.
    pub data: [u8; MAX_PAYLOAD_SIZE],
    /// Final magic number.
    pub magic_end: u32,
}

const _: () = {
    // Ensure block is correct size.
    assert!(core::mem::size_of::<Block>() == BLOCK_SIZE);
};

impl Default for Block {
    fn default() -> Self {
        Self {
            magic_start_0: MAGIC_NUMBER[0],
            magic_start_1: MAGIC_NUMBER[1],
            flags: Flags::default(),
            target_addr: 0,
            data_len: 0,
            block: 0,
            total_blocks: 0,
            board_family_id_or_file_size: 0,
            data: [0; MAX_PAYLOAD_SIZE],
            magic_end: MAGIC_NUMBER[2],
        }
    }
}

impl Block {
    pub fn new(
        block: usize,
        total_blocks: usize,
        data: &[u8],
        target_addr: usize,
    ) -> Self {
        // default with correct magic numbers
        let mut this = Self::default();

        // block index and total
        assert!(block <= total_blocks);
        assert!(block <= u32::MAX as usize);
        this.block = block as u32;
        assert!(total_blocks <= u32::MAX as usize);
        this.total_blocks = total_blocks as u32;

        // target flash address
        assert!(target_addr <= u32::MAX as usize);
        this.target_addr = target_addr as u32;

        // copy over data
        assert!(data.len() <= this.data.len());
        this.data_len = data.len() as u32;
        this.data[0..data.len()].copy_from_slice(data);

        this
    }

    /// Construct a [`Block`] from a slice.
    ///
    /// Returns an error if critical fields are incorrect.
    pub fn from_bytes(buf: &[u8]) -> Result<Block, BlockError> {
        let block = match Block::ref_from_bytes(buf) {
            Ok(b) => b,
            // INFO: e could be used for more detailed error
            Err(_e) => return Err(BlockError::InputBuffer),
        };

        if [block.magic_start_0, block.magic_start_1, block.magic_end]
            != MAGIC_NUMBER
        {
            return Err(BlockError::MagicNumber);
        }

        if block.data_len > MAX_PAYLOAD_SIZE as u32 {
            return Err(BlockError::PayloadSize);
        }

        // Validate block number is less than total blocks (UF2 spec requirement)
        if block.block >= block.total_blocks {
            return Err(BlockError::BlockNumberInvalid);
        }

        Ok(*block)
    }

    /// Returns `true` if the checksum flag is set.
    pub fn has_checksum(&self) -> bool {
        self.flags.contains(Flags::Checksum)
    }

    /// Returns the checksum value only if the checksum flag is set.
    pub fn checksum(&self) -> Option<&Checksum> {
        if self.has_checksum() {
            let len = self.data.len();
            Checksum::ref_from_bytes(&self.data[len - CHECKSUM_SIZE..len]).ok()
        } else {
            None
        }
    }

    /// Set the checksum for this block.
    pub fn set_checksum(&mut self, checksum: Checksum) {
        let begin = self.data.len() - size_of::<Checksum>();
        let end = self.data.len();

        self.data[begin..end].copy_from_slice(checksum.as_bytes());

        self.flags |= Flags::Checksum;
    }

    /// Returns `true` if the extensions flag is set.
    pub fn has_extensions(&self) -> bool {
        self.flags.contains(Flags::ExtensionTags)
    }

    /// Returns an extension [`Iterator`].
    pub fn extensions(&self) -> Option<Extensions<'_>> {
        if self.has_extensions() {
            let start = self.data_len as usize;
            let start = start.next_multiple_of(Extensions::ALIGN);
            let end = self.data.len();
            Some(Extensions::from_bytes(&self.data[start..end]))
        } else {
            None
        }
    }

    /// Returns the board family ID if the family ID flag is set.
    pub fn board_family_id(&self) -> Option<u32> {
        match self.flags.contains(Flags::FamilyId) {
            true => Some(self.board_family_id_or_file_size),
            false => None,
        }
    }

    /// Returns the payload data slice.
    pub fn data(&self) -> &[u8] {
        &self.data[0..self.data_len as usize]
    }

    /// Returns the file size if the family ID flag is not set.
    pub fn file_size(&self) -> Option<u32> {
        match self.flags.contains(Flags::FamilyId) {
            false => Some(self.board_family_id_or_file_size),
            true => None,
        }
    }

    /// Add an extension to this block.
    ///
    /// Extensions are stored after the payload data in the block's data array.
    /// Each extension consists of a 1-byte length field, a 3-byte tag, and the data.
    ///
    /// # Arguments
    /// * `tag` - The extension tag identifying the type of extension
    /// * `data` - The extension data bytes
    ///
    /// # Returns
    /// * `Ok(())` if the extension was added successfully
    /// * `Err(BlockError::PayloadSize)` if there's not enough space for the extension
    ///
    /// # Example
    /// ```
    /// use uftwo::block::{Block, ExtensionTag};
    ///
    /// let mut block = Block::new(0, 1, &[0xAA; 100], 0x08000000);
    /// block.add_extension(ExtensionTag::SemverString, b"1.0.0").unwrap();
    /// assert!(block.has_extensions());
    /// ```
    pub fn add_extension(
        &mut self,
        tag: ExtensionTag,
        data: &[u8],
    ) -> Result<(), BlockError> {
        // Find the end of existing extensions or start after payload
        let mut ext_start = self.data_len as usize;
        ext_start = ext_start.next_multiple_of(Extensions::ALIGN);

        // If there are existing extensions, find the end of the last one
        if self.has_extensions() {
            let existing_extensions = self.extensions().unwrap();
            for ext in existing_extensions {
                // Move to the end of this extension
                // The extension's total length is stored in the first byte
                let ext_total_len = Extensions::HEADER_SIZE + ext.data.len();
                ext_start += ext_total_len.next_multiple_of(Extensions::ALIGN);
            }
        }

        // Calculate space needed for this extension
        let ext_len = Extensions::HEADER_SIZE + data.len();
        let ext_end = ext_start + ext_len;

        // Check if there's enough space in the block's data array
        if ext_end > self.data.len() {
            return Err(BlockError::PayloadSize);
        }

        // Write extension length (1 byte)
        self.data[ext_start] = ext_len as u8;

        // Write extension tag (3 bytes)
        let tag_bytes = tag.to_bytes();
        self.data[ext_start + 1..ext_start + 4].copy_from_slice(&tag_bytes);

        // Write extension data
        self.data[ext_start + Extensions::HEADER_SIZE..ext_end]
            .copy_from_slice(data);

        // Set the extension flag
        self.flags |= Flags::ExtensionTags;

        Ok(())
    }
}

/// Checksum information.
///
/// This is used to allow skipping over blocks that do not need to be written
/// because the data has not changed.
#[derive(
    Debug, PartialEq, Eq, Immutable, KnownLayout, FromBytes, IntoBytes,
)]
#[repr(C)]
#[cfg_attr(feature = "defmt", derive(defmt::Format))]
pub struct Checksum {
    pub start: u32,
    pub length: u32,
    pub checksum: [u8; 16],
}

const _: () = {
    // Ensure Checksum is correct size.
    assert!(core::mem::size_of::<Checksum>() == CHECKSUM_SIZE);
};

/// Block flags.
#[derive(
    Debug, Default, Clone, Copy, PartialEq, Eq, Immutable, FromBytes, IntoBytes,
)]
#[repr(C)]
#[cfg_attr(feature = "defmt", derive(defmt::Format))]
pub struct Flags(u32);

bitflags::bitflags! {
    impl Flags: u32 {
        const NotMainFlash = 0x00000001;
        const FileContainer = 0x00001000;
        const FamilyId = 0x00002000;
        const Checksum = 0x00004000;
        const ExtensionTags = 0x00008000;
        const _ = !0; // non exhaustive
    }
}

/// Extensions access.
///
/// Use the `.next()` method to iterate through all of th extensions in the
/// current block. `.next()` will return `None` when there are no more
/// extensions left or none defined in the first place.
#[derive(Debug)]
#[cfg_attr(feature = "defmt", derive(defmt::Format))]
pub struct Extensions<'a> {
    start: usize,
    data: &'a [u8],
}

impl<'a> Extensions<'a> {
    /// Length byte + tag bytes.
    const HEADER_SIZE: usize = 4;

    /// Align to 4 byte boundary
    const ALIGN: usize = 4;

    /// Create a new extension iterator from bytes.
    pub fn from_bytes(data: &'a [u8]) -> Self {
        Self { start: 0, data }
    }

    fn current_tag(&self) -> ExtensionTag {
        let tag = u32::from_le_bytes([
            self.data[self.start + 1],
            self.data[self.start + 2],
            self.data[self.start + 3],
            0,
        ]);
        ExtensionTag::from(tag)
    }
}

impl<'a> Iterator for Extensions<'a> {
    type Item = Extension<'a>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.start > self.data.len() {
            // we are at the end
            return None;
        }

        let len = self.data[self.start] as usize;

        if self.start + Self::HEADER_SIZE > self.start + len {
            // there is no more tags
            return None;
        }

        let extension = Extension {
            tag: self.current_tag(),
            data: &self.data[self.start + Self::HEADER_SIZE..self.start + len],
        };

        // incerment start point
        // i.e where does the next (potential) tag start
        self.start += len;
        self.start = self.start.next_multiple_of(Self::ALIGN);

        Some(extension)
    }
}

/// An additional piece of information which can be appended after payload
/// data.
///
/// Converting the extension tag to UTF-8 strings or otherwise is an exercise
/// left to the user.
#[derive(Debug)]
#[cfg_attr(feature = "defmt", derive(defmt::Format))]
pub struct Extension<'a> {
    pub tag: ExtensionTag,
    pub data: &'a [u8],
}

impl<'a> Extension<'a> {
    /// Length byte + tag bytes.
    pub const HEADER_SIZE: usize = 4;
}

/// Extension tag.
#[derive(Debug, PartialEq, Eq)]
#[repr(u32)]
#[cfg_attr(feature = "defmt", derive(defmt::Format))]
pub enum ExtensionTag {
    /// UTF-8 Semantic Versioning string.
    SemverString = 0x9fc7bc,
    /// UTF-8 device description.
    DescriptionString = 0x650d9d,
    /// Page size of target device.
    TargetPageSize = 0x0be9f7,
    /// SHA-2 checksum of the firmware.
    Sha2Checksum = 0xb46db0,
    /// Device type identifier.
    DeviceTypeId = 0xc8a729,
    /// Other unknown tag.
    Other(u32),
}

impl From<u32> for ExtensionTag {
    fn from(value: u32) -> Self {
        match value {
            0x9fc7bc => Self::SemverString,
            0x650d9d => Self::DescriptionString,
            0x0be9f7 => Self::TargetPageSize,
            0xb46db0 => Self::Sha2Checksum,
            0xc8a729 => Self::DeviceTypeId,
            _ => Self::Other(value), // still valid, just unknown to us
        }
    }
}

impl ExtensionTag {
    /// Convert the extension tag to its 3-byte representation.
    pub fn to_bytes(&self) -> [u8; 3] {
        match self {
            ExtensionTag::SemverString => {
                0x9fc7bc_u32.to_le_bytes()[0..3].try_into().unwrap()
            }
            ExtensionTag::DescriptionString => {
                0x650d9d_u32.to_le_bytes()[0..3].try_into().unwrap()
            }
            ExtensionTag::TargetPageSize => {
                0x0be9f7_u32.to_le_bytes()[0..3].try_into().unwrap()
            }
            ExtensionTag::Sha2Checksum => {
                0xb46db0_u32.to_le_bytes()[0..3].try_into().unwrap()
            }
            ExtensionTag::DeviceTypeId => {
                0xc8a729_u32.to_le_bytes()[0..3].try_into().unwrap()
            }
            ExtensionTag::Other(value) => {
                value.to_le_bytes()[0..3].try_into().unwrap()
            }
        }
    }
}

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

    #[test]
    fn magic_number() {
        assert_eq!(MAGIC_NUMBER[0].as_bytes(), b"UF2\n");
    }

    #[test]
    fn block_checksum() {
        let mut block = Block::default();
        assert_eq!(block.has_checksum(), false);

        block.flags |= Flags::Checksum;
        assert_eq!(block.has_checksum(), true);

        let cksm = block.checksum();
        assert!(cksm.is_some());
    }

    #[test]
    fn block_extension() {
        let mut block = Block {
            flags: Flags::ExtensionTags,
            data_len: 0,
            ..Default::default()
        };

        // Semver string
        block.data[0..12].copy_from_slice(&[
            0x09, 0xbc, 0xc7, 0x9f, 0x30, 0x2e, 0x31, 0x2e, 0x32, 0x00, 0x00,
            0x00,
        ]);
        // Semver string
        block.data[12..24].copy_from_slice(&[
            0x09, 0xbc, 0xc7, 0x9f, 0x30, 0x2e, 0x31, 0x2e, 0x32, 0x00, 0x00,
            0x00,
        ]);
        // Device description
        block.data[24..44].copy_from_slice(&[
            0x14, 0x9d, 0x0d, 0x65, 0x41, 0x43, 0x4d, 0x45, 0x20, 0x54, 0x6f,
            0x61, 0x73, 0x74, 0x65, 0x72, 0x20, 0x6d, 0x6b, 0x33,
        ]);

        assert!(block.extensions().is_some());

        let mut extensions = block.extensions().unwrap();

        let first = extensions.next().unwrap();
        assert_eq!(first.tag, ExtensionTag::SemverString);
        assert_eq!(first.data, b"0.1.2");

        let second = extensions.next().unwrap();
        assert_eq!(second.tag, ExtensionTag::SemverString);
        assert_eq!(second.data, b"0.1.2");

        let third = extensions.next().unwrap();
        assert_eq!(third.tag, ExtensionTag::DescriptionString);
        assert_eq!(third.data, b"ACME Toaster mk3");
    }

    #[test]
    fn example_file() {
        use std::io::prelude::*;

        let mut f = std::fs::File::open("example.uf2").unwrap();
        let mut buffer = [0; 512];

        f.read(&mut buffer).unwrap();

        let block = Block::from_bytes(&buffer).unwrap();

        assert_eq!(block.magic_start_0, MAGIC_NUMBER[0]);
        assert_eq!(block.magic_start_1, MAGIC_NUMBER[1]);
        assert_eq!(block.magic_end, MAGIC_NUMBER[2]);

        assert_eq!(block.target_addr, 0x2000);
        assert_eq!(block.data_len, 256);
        assert_eq!(block.block, 0);
        assert_eq!(block.total_blocks, 1438);
        assert_eq!(block.board_family_id_or_file_size, 0);
    }

    #[test]
    fn test_block_new() {
        // Test basic block creation
        let data = [0xAA; 256];
        let block = Block::new(0, 1, &data, 0x08000000);

        // Check magic numbers
        assert_eq!(block.magic_start_0, MAGIC_NUMBER[0]);
        assert_eq!(block.magic_start_1, MAGIC_NUMBER[1]);
        assert_eq!(block.magic_end, MAGIC_NUMBER[2]);

        // Check metadata
        assert_eq!(block.block, 0);
        assert_eq!(block.total_blocks, 1);
        assert_eq!(block.target_addr, 0x08000000);
        assert_eq!(block.data_len, 256);

        // Check data
        assert_eq!(block.data(), &data);
    }

    #[test]
    fn test_block_new_multiple_blocks() {
        // Test with multiple blocks
        let data = [0xBB; 100];
        let block = Block::new(2, 5, &data, 0x08000100);

        assert_eq!(block.block, 2);
        assert_eq!(block.total_blocks, 5);
        assert_eq!(block.target_addr, 0x08000100);
        assert_eq!(block.data_len, 100);
        assert_eq!(block.data(), &data);
    }

    #[test]
    fn test_block_new_empty_data() {
        // Test with empty data
        let data: &[u8] = &[];
        let block = Block::new(0, 1, data, 0);

        assert_eq!(block.data_len, 0);
        assert_eq!(block.data(), &[]);
    }

    #[test]
    fn test_block_new_max_payload() {
        // Test with maximum payload size
        let data = [0xCC; MAX_PAYLOAD_SIZE];
        let block = Block::new(0, 1, &data, 0);

        assert_eq!(block.data_len, MAX_PAYLOAD_SIZE as u32);
        assert_eq!(block.data(), &data);
    }

    #[test]
    #[should_panic(expected = "block <= total_blocks")]
    fn test_block_new_panics_on_invalid_index() {
        // Block index cannot exceed total blocks
        let data = [0xDD; 100];
        Block::new(5, 3, &data, 0); // block=5 > total_blocks=3
    }

    #[test]
    fn test_from_bytes_block_number_exceeds_total() {
        // Create a block with block >= total_blocks
        let mut block = Block::default();
        block.block = 5; // block number 5
        block.total_blocks = 3; // but only 3 total blocks
        let bytes = block.as_bytes();

        let result = Block::from_bytes(&bytes);
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            BlockError::BlockNumberInvalid
        ));
    }

    #[test]
    fn test_set_checksum() {
        let mut block = Block::default();

        // Initially no checksum
        assert_eq!(block.has_checksum(), false);
        assert!(block.checksum().is_none());

        // Set a checksum
        let checksum = Checksum {
            start: 0x08000000,
            length: 256,
            checksum: [0xAB; 16],
        };
        block.set_checksum(checksum);

        // Verify checksum is set
        assert_eq!(block.has_checksum(), true);

        // Verify we can retrieve the checksum
        let retrieved = block.checksum().unwrap();
        assert_eq!(retrieved.start, 0x08000000);
        assert_eq!(retrieved.length, 256);
        assert_eq!(retrieved.checksum, [0xAB; 16]);
    }

    #[test]
    fn test_board_family_id() {
        let mut block = Block::default();

        // Initially no family ID
        assert_eq!(block.board_family_id(), None);

        // Set family ID flag and value
        block.flags |= Flags::FamilyId;
        block.board_family_id_or_file_size = 0x12345678;

        // Verify we can retrieve the family ID
        assert_eq!(block.board_family_id(), Some(0x12345678));

        // Test with different family ID
        block.board_family_id_or_file_size = 0x87654321;
        assert_eq!(block.board_family_id(), Some(0x87654321));
    }

    #[test]
    fn test_file_size() {
        let mut block = Block::default();

        // Initially no file size (FamilyId flag not set)
        assert_eq!(block.file_size(), Some(0));

        // Set a file size
        block.board_family_id_or_file_size = 1024;
        assert_eq!(block.file_size(), Some(1024));

        // Set FamilyId flag - now file_size should return None
        block.flags |= Flags::FamilyId;
        assert_eq!(block.file_size(), None);

        // Clear FamilyId flag - file_size should return the value again
        block.flags &= !Flags::FamilyId;
        assert_eq!(block.file_size(), Some(1024));
    }

    #[test]
    fn test_board_family_id_vs_file_size() {
        let mut block = Block::default();

        // Set both flags and value
        block.board_family_id_or_file_size = 0xCAFEBABE;

        // Without FamilyId flag, it's a file size
        assert_eq!(block.file_size(), Some(0xCAFEBABE));
        assert_eq!(block.board_family_id(), None);

        // With FamilyId flag, it's a family ID
        block.flags |= Flags::FamilyId;
        assert_eq!(block.board_family_id(), Some(0xCAFEBABE));
        assert_eq!(block.file_size(), None);

        // The same underlying value, but different interpretation based on flag
        assert_eq!(block.board_family_id_or_file_size, 0xCAFEBABE);
    }

    #[test]
    fn test_add_extension() {
        let mut block = Block::new(0, 1, &[0xAA; 100], 0x08000000);

        // Initially no extensions
        assert_eq!(block.has_extensions(), false);

        // Add a semver extension
        let result = block.add_extension(ExtensionTag::SemverString, b"1.0.0");
        assert!(result.is_ok());
        assert_eq!(block.has_extensions(), true);

        // Verify we can retrieve the extension
        let mut extensions = block.extensions().unwrap();
        let ext = extensions.next().unwrap();
        assert_eq!(ext.tag, ExtensionTag::SemverString);
        assert_eq!(ext.data, b"1.0.0");
    }

    #[test]
    fn test_add_extension_multiple() {
        let mut block = Block::new(0, 1, &[0xBB; 50], 0x08000000);

        // Add first extension
        block
            .add_extension(ExtensionTag::SemverString, b"1.0.0")
            .unwrap();

        // Verify first extension is present
        let mut extensions = block.extensions().unwrap();
        let ext1 = extensions.next().unwrap();
        assert_eq!(ext1.tag, ExtensionTag::SemverString);
        assert_eq!(ext1.data, b"1.0.0");

        // Add second extension
        block
            .add_extension(ExtensionTag::DescriptionString, b"Test")
            .unwrap();

        // Verify both extensions are present
        let mut extensions = block.extensions().unwrap();
        let ext1 = extensions.next().unwrap();
        let ext2 = extensions.next().unwrap();

        // Check that we have both extensions
        let tag1 = ext1.tag;
        let tag2 = ext2.tag;
        assert_eq!(tag1, ExtensionTag::SemverString);
        assert_eq!(tag2, ExtensionTag::DescriptionString);
    }

    #[test]
    fn test_add_extension_no_space() {
        // Create a block with maximum payload to leave no space for extensions
        let mut block = Block::new(0, 1, &[0xCC; MAX_PAYLOAD_SIZE], 0x08000000);

        // Try to add an extension when there's no space left
        let result = block.add_extension(ExtensionTag::SemverString, b"1.0.0");
        assert!(result.is_err());
        assert!(matches!(result.unwrap_err(), BlockError::PayloadSize));
    }
}