//! # Partition Table Support
//!
//! ## Overview
//!
//! This module allows reading the partition table and conveniently
//! writing/reading partition contents.
//!
//! For more information see <https://docs.espressif.com/projects/esp-idf/en/latest/esp32/api-guides/partition-tables.html#built-in-partition-tables>

/// Maximum length of a partition table.
pub const PARTITION_TABLE_MAX_LEN: usize = 0xC00;

const PARTITION_TABLE_OFFSET: u32 =
    esp_config::esp_config_int!(u32, "ESP_BOOTLOADER_ESP_IDF_CONFIG_PARTITION_TABLE_OFFSET");

const RAW_ENTRY_LEN: usize = 32;
const ENTRY_MAGIC: u16 = 0x50aa;
const MD5_MAGIC: u16 = 0xebeb;

const OTA_SUBTYPE_OFFSET: u8 = 0x10;

/// Represents a single partition entry.
#[derive(Clone, Copy)]
pub struct PartitionEntry<'a> {
    pub(crate) binary: &'a [u8; RAW_ENTRY_LEN],
}

impl<'a> PartitionEntry<'a> {
    fn new(binary: &'a [u8; RAW_ENTRY_LEN]) -> Self {
        Self { binary }
    }

    /// The magic value of the entry.
    pub fn magic(&self) -> u16 {
        u16::from_le_bytes(unwrap!(self.binary[..2].try_into()))
    }

    /// The partition type in raw representation.
    pub fn raw_type(&self) -> u8 {
        self.binary[2]
    }

    /// The partition sub-type in raw representation.
    pub fn raw_subtype(&self) -> u8 {
        self.binary[3]
    }

    /// Offset of the partition on flash.
    pub fn offset(&self) -> u32 {
        u32::from_le_bytes(unwrap!(self.binary[4..][..4].try_into()))
    }

    /// Length of the partition in bytes.
    pub fn len(&self) -> u32 {
        u32::from_le_bytes(unwrap!(self.binary[8..][..4].try_into()))
    }

    /// Checks for a zero-length partition.
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// The label of the partition.
    pub fn label(&self) -> &'a [u8] {
        &self.binary[12..][..16]
    }

    /// The label of the partition as `&str`.
    pub fn label_as_str(&self) -> &'a str {
        let array = self.label();
        let len = array
            .iter()
            .position(|b| *b == 0 || *b == 0xff)
            .unwrap_or(array.len());
        unsafe {
            core::str::from_utf8_unchecked(core::slice::from_raw_parts(array.as_ptr().cast(), len))
        }
    }

    /// Raw flags of this partition. You probably want to use
    /// [Self::is_read_only] and [Self::is_encrypted] instead.
    pub fn flags(&self) -> u32 {
        u32::from_le_bytes(unwrap!(self.binary[28..][..4].try_into()))
    }

    /// If the partition is read only.
    pub fn is_read_only(&self) -> bool {
        self.flags() & 0b01 != 0
    }

    /// If the partition is encrypted.
    pub fn is_encrypted(&self) -> bool {
        self.flags() & 0b10 != 0
    }

    /// The partition type (type and sub-type).
    pub fn partition_type(&self) -> PartitionType {
        match self.raw_type() {
            0 => PartitionType::App(unwrap!(self.raw_subtype().try_into())),
            1 => PartitionType::Data(unwrap!(self.raw_subtype().try_into())),
            2 => PartitionType::Bootloader(unwrap!(self.raw_subtype().try_into())),
            3 => PartitionType::PartitionTable(unwrap!(self.raw_subtype().try_into())),
            _ => unreachable!(),
        }
    }

    /// Provides a "view" into the partition allowing to read/write the
    /// partition contents by using the given [embedded_storage::Storage] and/or
    /// [embedded_storage::ReadStorage] implementation.
    pub fn as_embedded_storage<F>(self, flash: &'a mut F) -> FlashRegion<'a, F>
    where
        F: embedded_storage::ReadStorage,
    {
        FlashRegion { raw: self, flash }
    }
}

impl core::fmt::Debug for PartitionEntry<'_> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("PartitionEntry")
            .field("magic", &self.magic())
            .field("raw_type", &self.raw_type())
            .field("raw_subtype", &self.raw_subtype())
            .field("offset", &self.offset())
            .field("len", &self.len())
            .field("label", &self.label_as_str())
            .field("flags", &self.flags())
            .field("is_read_only", &self.is_read_only())
            .field("is_encrypted", &self.is_encrypted())
            .finish()
    }
}

#[cfg(feature = "defmt")]
impl defmt::Format for PartitionEntry<'_> {
    fn format(&self, fmt: defmt::Formatter) {
        defmt::write!(
            fmt,
            "PartitionEntry (\
            magic = {}, \
            raw_type = {}, \
            raw_subtype = {}, \
            offset = {}, \
            len = {}, \
            label = {}, \
            flags = {}, \
            is_read_only = {}, \
            is_encrypted = {}\
            )",
            self.magic(),
            self.raw_type(),
            self.raw_subtype(),
            self.offset(),
            self.len(),
            self.label_as_str(),
            self.flags(),
            self.is_read_only(),
            self.is_encrypted()
        )
    }
}

/// Errors which can be returned.
#[derive(Debug, PartialEq, Eq, Clone, Copy, Hash, strum::Display)]
#[cfg_attr(feature = "defmt", derive(defmt::Format))]
#[non_exhaustive]
pub enum Error {
    /// The partition table is invalid or doesn't contain a needed partition.
    Invalid,
    /// An operation tries to access data that is out of bounds.
    OutOfBounds,
    /// An error which originates from the embedded-storage implementation.
    StorageError,
    /// The partition is write protected.
    WriteProtected,
    /// The partition is invalid.
    InvalidPartition {
        expected_size: usize,
        expected_type: PartitionType,
    },
    /// Invalid tate
    InvalidState,
    /// The given argument is invalid.
    InvalidArgument,
}

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

/// A partition table.
#[derive(Debug)]
#[cfg_attr(feature = "defmt", derive(defmt::Format))]
pub struct PartitionTable<'a> {
    binary: &'a [[u8; RAW_ENTRY_LEN]],
    entries: usize,
}

impl<'a> PartitionTable<'a> {
    fn new(binary: &'a [u8]) -> Result<Self, Error> {
        if binary.len() > PARTITION_TABLE_MAX_LEN {
            return Err(Error::Invalid);
        }

        let (binary, rem) = binary.as_chunks::<RAW_ENTRY_LEN>();
        if !rem.is_empty() {
            return Err(Error::Invalid);
        }

        if binary.is_empty() {
            return Ok(Self {
                binary: &[],
                entries: 0,
            });
        }

        let mut raw_table = Self {
            binary,
            entries: binary.len(),
        };

        #[cfg(feature = "validation")]
        {
            let (hash, index) = {
                let mut i = 0;
                loop {
                    if let Ok(entry) = raw_table.get_partition(i) {
                        if entry.magic() == MD5_MAGIC {
                            break (&entry.binary[16..][..16], i);
                        }

                        i += 1;
                        if i >= raw_table.entries {
                            return Err(Error::Invalid);
                        }
                    }
                }
            };

            let mut hasher = crate::crypto::Md5::new();

            for i in 0..index {
                hasher.update(&raw_table.binary[i]);
            }
            let calculated_hash = hasher.finalize();

            if calculated_hash != hash {
                return Err(Error::Invalid);
            }
        }

        let entries = {
            let mut i = 0;
            loop {
                if let Ok(entry) = raw_table.get_partition(i) {
                    if entry.magic() != ENTRY_MAGIC {
                        break;
                    }

                    i += 1;

                    if i == raw_table.entries {
                        break;
                    }
                } else {
                    return Err(Error::Invalid);
                }
            }
            i
        };

        raw_table.entries = entries;

        Ok(raw_table)
    }

    /// Number of partitions contained in the partition table.
    pub fn len(&self) -> usize {
        self.entries
    }

    /// Checks if there are no recognized partitions.
    pub fn is_empty(&self) -> bool {
        self.entries == 0
    }

    /// Get a partition entry.
    pub fn get_partition(&self, index: usize) -> Result<PartitionEntry<'a>, Error> {
        if index >= self.entries {
            return Err(Error::OutOfBounds);
        }
        Ok(PartitionEntry::new(&self.binary[index]))
    }

    /// Get the first partition matching the given partition type.
    pub fn find_partition(&self, pt: PartitionType) -> Result<Option<PartitionEntry<'a>>, Error> {
        for i in 0..self.entries {
            let entry = self.get_partition(i)?;
            if entry.partition_type() == pt {
                return Ok(Some(entry));
            }
        }
        Ok(None)
    }

    /// Returns an iterator over the partitions.
    pub fn iter(&self) -> impl Iterator<Item = PartitionEntry<'a>> {
        (0..self.entries).filter_map(|i| self.get_partition(i).ok())
    }

    #[cfg(feature = "std")]
    /// Get the currently booted partition.
    pub fn booted_partition(&self) -> Result<Option<PartitionEntry<'a>>, Error> {
        Err(Error::Invalid)
    }

    #[cfg(not(feature = "std"))]
    /// Get the currently booted partition.
    pub fn booted_partition(&self) -> Result<Option<PartitionEntry<'a>>, Error> {
        // Read entry 0 from MMU to know which partition is mapped
        //
        // See <https://github.com/espressif/esp-idf/blob/758939caecb16e5542b3adfba0bc85025517db45/components/hal/mmu_hal.c#L124>
        cfg_if::cfg_if! {
            if #[cfg(feature = "esp32")] {
                let paddr = unsafe {
                    ((0x3FF10000 as *const u32).read_volatile() & 0xff) << 16
                };
            } else if #[cfg(feature = "esp32s2")] {
                let paddr = unsafe {
                    (((0x61801000 + 128 * 4) as *const u32).read_volatile() & 0xff) << 16
                };
            } else if #[cfg(feature = "esp32s3")] {
                // Revisit this once we support XiP from PSRAM for ESP32-S3
                let paddr = unsafe {
                    ((0x600C5000 as *const u32).read_volatile() & 0xff) << 16
                };
            } else if #[cfg(any(feature = "esp32c2", feature = "esp32c3"))] {
                let paddr = unsafe {
                    ((0x600c5000 as *const u32).read_volatile() & 0xff) << 16
                };
            } else if #[cfg(any(feature = "esp32c6", feature = "esp32h2"))] {
                let paddr = unsafe {
                    ((0x60002000 + 0x380) as *mut u32).write_volatile(0);
                    (((0x60002000 + 0x37c) as *const u32).read_volatile() & 0xff) << 16
                };
            }
        }

        for id in 0..self.len() {
            let entry = self.get_partition(id)?;
            if entry.offset() == paddr {
                return Ok(Some(entry));
            }
        }

        Ok(None)
    }
}

/// A partition type including the sub-type.
#[derive(Debug, PartialEq, Eq, Clone, Copy, Hash)]
#[cfg_attr(feature = "defmt", derive(defmt::Format))]
pub enum PartitionType {
    /// Application.
    App(AppPartitionSubType),
    /// Data.
    Data(DataPartitionSubType),
    /// Bootloader.
    Bootloader(BootloaderPartitionSubType),
    /// Partition table.
    PartitionTable(PartitionTablePartitionSubType),
}

/// A partition type
#[derive(Debug, PartialEq, Eq, Clone, Copy, Hash)]
#[cfg_attr(feature = "defmt", derive(defmt::Format))]
#[repr(u8)]
pub enum RawPartitionType {
    /// Application.
    App = 0,
    /// Data.
    Data,
    /// Bootloader.
    Bootloader,
    /// Partition table.
    PartitionTable,
}

/// Sub-types of an application partition.
#[derive(Debug, PartialEq, Eq, Clone, Copy, Hash, strum::FromRepr)]
#[cfg_attr(feature = "defmt", derive(defmt::Format))]
#[repr(u8)]
pub enum AppPartitionSubType {
    /// Factory image
    Factory = 0,
    /// OTA slot 0
    Ota0    = OTA_SUBTYPE_OFFSET,
    /// OTA slot 1
    Ota1,
    /// OTA slot 2
    Ota2,
    /// OTA slot 3
    Ota3,
    /// OTA slot 4
    Ota4,
    /// OTA slot 5
    Ota5,
    /// OTA slot 6
    Ota6,
    /// OTA slot 7
    Ota7,
    /// OTA slot 8
    Ota8,
    /// OTA slot 9
    Ota9,
    /// OTA slot 10
    Ota10,
    /// OTA slot 11
    Ota11,
    /// OTA slot 12
    Ota12,
    /// OTA slot 13
    Ota13,
    /// OTA slot 14
    Ota14,
    /// OTA slot 15
    Ota15,
    /// Test image
    Test,
}

impl AppPartitionSubType {
    pub(crate) fn ota_app_number(&self) -> u8 {
        *self as u8 - OTA_SUBTYPE_OFFSET
    }

    pub(crate) fn from_ota_app_number(number: u8) -> Result<Self, Error> {
        if number > 16 {
            return Err(Error::InvalidArgument);
        }
        Self::try_from(number + OTA_SUBTYPE_OFFSET)
    }
}

impl TryFrom<u8> for AppPartitionSubType {
    type Error = Error;

    fn try_from(value: u8) -> Result<Self, Self::Error> {
        AppPartitionSubType::from_repr(value).ok_or(Error::Invalid)
    }
}

/// Sub-types of the data partition type.
#[derive(Debug, PartialEq, Eq, Clone, Copy, Hash, strum::FromRepr)]
#[cfg_attr(feature = "defmt", derive(defmt::Format))]
#[repr(u8)]
pub enum DataPartitionSubType {
    /// Data partition which stores information about the currently selected OTA
    /// app slot. This partition should be 0x2000 bytes in size. Refer to
    /// the OTA documentation for more details.
    Ota      = 0,
    /// Phy is for storing PHY initialization data. This allows PHY to be
    /// configured per-device, instead of in firmware.
    Phy,
    /// Used for Non-Volatile Storage (NVS).
    Nvs,
    /// Used for storing core dumps while using a custom partition table
    Coredump,
    /// NvsKeys is used for the NVS key partition. (NVS).
    NvsKeys,
    /// Used for emulating eFuse bits using Virtual eFuses.
    EfuseEm,
    /// Implicitly used for data partitions with unspecified (empty) subtype,
    /// but it is possible to explicitly mark them as undefined as well.
    Undefined,
    /// FAT Filesystem Support.
    Fat      = 0x81,
    /// SPIFFS Filesystem.
    Spiffs   = 0x82,
    ///  LittleFS filesystem.
    LittleFs = 0x83,
}

impl TryFrom<u8> for DataPartitionSubType {
    type Error = Error;

    fn try_from(value: u8) -> Result<Self, Self::Error> {
        DataPartitionSubType::from_repr(value).ok_or(Error::Invalid)
    }
}

/// Sub-type of the bootloader partition type.
#[derive(Debug, PartialEq, Eq, Clone, Copy, Hash, strum::FromRepr)]
#[cfg_attr(feature = "defmt", derive(defmt::Format))]
#[repr(u8)]
pub enum BootloaderPartitionSubType {
    /// It is the so-called 2nd stage bootloader.
    Primary = 0,
    /// It is a temporary bootloader partition used by the bootloader OTA update
    /// functionality for downloading a new image.
    Ota     = 1,
}

impl TryFrom<u8> for BootloaderPartitionSubType {
    type Error = Error;

    fn try_from(value: u8) -> Result<Self, Self::Error> {
        BootloaderPartitionSubType::from_repr(value).ok_or(Error::Invalid)
    }
}

/// Sub-type of the partition table type.
#[derive(Debug, PartialEq, Eq, Clone, Copy, Hash, strum::FromRepr)]
#[cfg_attr(feature = "defmt", derive(defmt::Format))]
#[repr(u8)]
pub enum PartitionTablePartitionSubType {
    /// It is the primary partition table.
    Primary = 0,
    /// It is a temporary partition table partition used by the partition table
    /// OTA update functionality for downloading a new image.
    Ota     = 1,
}

impl TryFrom<u8> for PartitionTablePartitionSubType {
    type Error = Error;

    fn try_from(value: u8) -> Result<Self, Self::Error> {
        PartitionTablePartitionSubType::from_repr(value).ok_or(Error::Invalid)
    }
}

/// Read the partition table.
///
/// Pass an implementation of [embedded_storage::Storage] which can read from
/// the whole flash and provide storage to read the partition table into.
pub fn read_partition_table<'a>(
    flash: &mut impl embedded_storage::Storage,
    storage: &'a mut [u8],
) -> Result<PartitionTable<'a>, Error> {
    flash
        .read(PARTITION_TABLE_OFFSET, storage)
        .map_err(|_e| Error::StorageError)?;

    PartitionTable::new(storage)
}

/// A flash region is a "view" into the partition.
///
/// It allows to read and write to the partition without the need to account for
/// the partition offset.
#[derive(Debug)]
#[cfg_attr(feature = "defmt", derive(defmt::Format))]
pub struct FlashRegion<'a, F> {
    pub(crate) raw: PartitionEntry<'a>,
    pub(crate) flash: &'a mut F,
}

impl<F> FlashRegion<'_, F> {
    /// Returns the size of the partition in bytes.
    pub fn partition_size(&self) -> usize {
        self.raw.len() as _
    }

    fn range(&self) -> core::ops::Range<u32> {
        self.raw.offset()..self.raw.offset() + self.raw.len()
    }

    fn in_range(&self, start: u32, len: usize) -> bool {
        self.range().contains(&start) && (start + len as u32 <= self.range().end)
    }
}

impl<F> embedded_storage::Region for FlashRegion<'_, F> {
    fn contains(&self, address: u32) -> bool {
        self.range().contains(&address)
    }
}

impl<F> embedded_storage::ReadStorage for FlashRegion<'_, F>
where
    F: embedded_storage::ReadStorage,
{
    type Error = Error;

    fn read(&mut self, offset: u32, bytes: &mut [u8]) -> Result<(), Self::Error> {
        let address = offset + self.raw.offset();

        if !self.in_range(address, bytes.len()) {
            return Err(Error::OutOfBounds);
        }

        self.flash
            .read(address, bytes)
            .map_err(|_e| Error::StorageError)
    }

    fn capacity(&self) -> usize {
        self.partition_size()
    }
}

impl<F> embedded_storage::Storage for FlashRegion<'_, F>
where
    F: embedded_storage::Storage,
{
    fn write(&mut self, offset: u32, bytes: &[u8]) -> Result<(), Self::Error> {
        let address = offset + self.raw.offset();

        if self.raw.is_read_only() {
            return Err(Error::WriteProtected);
        }

        if !self.in_range(address, bytes.len()) {
            return Err(Error::OutOfBounds);
        }

        self.flash
            .write(address, bytes)
            .map_err(|_e| Error::StorageError)
    }
}

impl embedded_storage::nor_flash::NorFlashError for Error {
    fn kind(&self) -> embedded_storage::nor_flash::NorFlashErrorKind {
        match self {
            Error::OutOfBounds => embedded_storage::nor_flash::NorFlashErrorKind::OutOfBounds,
            _ => embedded_storage::nor_flash::NorFlashErrorKind::Other,
        }
    }
}

impl<F> embedded_storage::nor_flash::ErrorType for FlashRegion<'_, F> {
    type Error = Error;
}

impl<F> embedded_storage::nor_flash::ReadNorFlash for FlashRegion<'_, F>
where
    F: embedded_storage::nor_flash::ReadNorFlash,
{
    const READ_SIZE: usize = F::READ_SIZE;

    fn read(&mut self, offset: u32, bytes: &mut [u8]) -> Result<(), Self::Error> {
        let address = offset + self.raw.offset();

        if !self.in_range(address, bytes.len()) {
            return Err(Error::OutOfBounds);
        }

        self.flash
            .read(address, bytes)
            .map_err(|_e| Error::StorageError)
    }

    fn capacity(&self) -> usize {
        self.partition_size()
    }
}

impl<F> embedded_storage::nor_flash::NorFlash for FlashRegion<'_, F>
where
    F: embedded_storage::nor_flash::NorFlash,
{
    const WRITE_SIZE: usize = F::WRITE_SIZE;

    const ERASE_SIZE: usize = F::ERASE_SIZE;

    fn erase(&mut self, from: u32, to: u32) -> Result<(), Self::Error> {
        let address_from = from + self.raw.offset();
        let address_to = to + self.raw.offset();

        if self.raw.is_read_only() {
            return Err(Error::WriteProtected);
        }

        if !self.range().contains(&address_from) {
            return Err(Error::OutOfBounds);
        }

        if !self.range().contains(&address_to) {
            return Err(Error::OutOfBounds);
        }

        self.flash
            .erase(address_from, address_to)
            .map_err(|_e| Error::StorageError)
    }

    fn write(&mut self, offset: u32, bytes: &[u8]) -> Result<(), Self::Error> {
        let address = offset + self.raw.offset();

        if self.raw.is_read_only() {
            return Err(Error::WriteProtected);
        }

        if !self.in_range(address, bytes.len()) {
            return Err(Error::OutOfBounds);
        }

        self.flash
            .write(address, bytes)
            .map_err(|_e| Error::StorageError)
    }
}

impl<F> embedded_storage::nor_flash::MultiwriteNorFlash for FlashRegion<'_, F> where
    F: embedded_storage::nor_flash::MultiwriteNorFlash
{
}

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

    static SIMPLE: &[u8] = include_bytes!("../testdata/single_factory_no_ota.bin");
    static OTA: &[u8] = include_bytes!("../testdata/factory_app_two_ota.bin");

    #[test]
    fn read_simple() {
        let pt = PartitionTable::new(SIMPLE).unwrap();

        assert_eq!(3, pt.len());

        assert_eq!(1, pt.get_partition(0).unwrap().raw_type());
        assert_eq!(1, pt.get_partition(1).unwrap().raw_type());
        assert_eq!(0, pt.get_partition(2).unwrap().raw_type());

        assert_eq!(2, pt.get_partition(0).unwrap().raw_subtype());
        assert_eq!(1, pt.get_partition(1).unwrap().raw_subtype());
        assert_eq!(0, pt.get_partition(2).unwrap().raw_subtype());

        assert_eq!(
            PartitionType::Data(DataPartitionSubType::Nvs),
            pt.get_partition(0).unwrap().partition_type()
        );
        assert_eq!(
            PartitionType::Data(DataPartitionSubType::Phy),
            pt.get_partition(1).unwrap().partition_type()
        );
        assert_eq!(
            PartitionType::App(AppPartitionSubType::Factory),
            pt.get_partition(2).unwrap().partition_type()
        );

        assert_eq!(0x9000, pt.get_partition(0).unwrap().offset());
        assert_eq!(0xf000, pt.get_partition(1).unwrap().offset());
        assert_eq!(0x10000, pt.get_partition(2).unwrap().offset());

        assert_eq!(0x6000, pt.get_partition(0).unwrap().len());
        assert_eq!(0x1000, pt.get_partition(1).unwrap().len());
        assert_eq!(0x100000, pt.get_partition(2).unwrap().len());

        assert_eq!("nvs", pt.get_partition(0).unwrap().label_as_str());
        assert_eq!("phy_init", pt.get_partition(1).unwrap().label_as_str());
        assert_eq!("factory", pt.get_partition(2).unwrap().label_as_str());

        assert_eq!(false, pt.get_partition(0).unwrap().is_read_only());
        assert_eq!(false, pt.get_partition(1).unwrap().is_read_only());
        assert_eq!(false, pt.get_partition(2).unwrap().is_read_only());

        assert_eq!(false, pt.get_partition(0).unwrap().is_encrypted());
        assert_eq!(false, pt.get_partition(1).unwrap().is_encrypted());
        assert_eq!(false, pt.get_partition(2).unwrap().is_encrypted());
    }

    #[test]
    fn read_ota() {
        let pt = PartitionTable::new(OTA).unwrap();

        assert_eq!(6, pt.len());

        assert_eq!(1, pt.get_partition(0).unwrap().raw_type());
        assert_eq!(1, pt.get_partition(1).unwrap().raw_type());
        assert_eq!(1, pt.get_partition(2).unwrap().raw_type());
        assert_eq!(0, pt.get_partition(3).unwrap().raw_type());
        assert_eq!(0, pt.get_partition(4).unwrap().raw_type());
        assert_eq!(0, pt.get_partition(5).unwrap().raw_type());

        assert_eq!(2, pt.get_partition(0).unwrap().raw_subtype());
        assert_eq!(0, pt.get_partition(1).unwrap().raw_subtype());
        assert_eq!(1, pt.get_partition(2).unwrap().raw_subtype());
        assert_eq!(0, pt.get_partition(3).unwrap().raw_subtype());
        assert_eq!(0x10, pt.get_partition(4).unwrap().raw_subtype());
        assert_eq!(0x11, pt.get_partition(5).unwrap().raw_subtype());

        assert_eq!(
            PartitionType::Data(DataPartitionSubType::Nvs),
            pt.get_partition(0).unwrap().partition_type()
        );
        assert_eq!(
            PartitionType::Data(DataPartitionSubType::Ota),
            pt.get_partition(1).unwrap().partition_type()
        );
        assert_eq!(
            PartitionType::Data(DataPartitionSubType::Phy),
            pt.get_partition(2).unwrap().partition_type()
        );
        assert_eq!(
            PartitionType::App(AppPartitionSubType::Factory),
            pt.get_partition(3).unwrap().partition_type()
        );
        assert_eq!(
            PartitionType::App(AppPartitionSubType::Ota0),
            pt.get_partition(4).unwrap().partition_type()
        );
        assert_eq!(
            PartitionType::App(AppPartitionSubType::Ota1),
            pt.get_partition(5).unwrap().partition_type()
        );

        assert_eq!(0x9000, pt.get_partition(0).unwrap().offset());
        assert_eq!(0xd000, pt.get_partition(1).unwrap().offset());
        assert_eq!(0xf000, pt.get_partition(2).unwrap().offset());
        assert_eq!(0x10000, pt.get_partition(3).unwrap().offset());
        assert_eq!(0x110000, pt.get_partition(4).unwrap().offset());
        assert_eq!(0x210000, pt.get_partition(5).unwrap().offset());

        assert_eq!(0x4000, pt.get_partition(0).unwrap().len());
        assert_eq!(0x2000, pt.get_partition(1).unwrap().len());
        assert_eq!(0x1000, pt.get_partition(2).unwrap().len());
        assert_eq!(0x100000, pt.get_partition(3).unwrap().len());
        assert_eq!(0x100000, pt.get_partition(4).unwrap().len());
        assert_eq!(0x100000, pt.get_partition(5).unwrap().len());

        assert_eq!("nvs", pt.get_partition(0).unwrap().label_as_str());
        assert_eq!("otadata", pt.get_partition(1).unwrap().label_as_str());
        assert_eq!("phy_init", pt.get_partition(2).unwrap().label_as_str());
        assert_eq!("factory", pt.get_partition(3).unwrap().label_as_str());
        assert_eq!("ota_0", pt.get_partition(4).unwrap().label_as_str());
        assert_eq!("ota_1", pt.get_partition(5).unwrap().label_as_str());

        assert_eq!(false, pt.get_partition(0).unwrap().is_read_only());
        assert_eq!(false, pt.get_partition(1).unwrap().is_read_only());
        assert_eq!(false, pt.get_partition(2).unwrap().is_read_only());
        assert_eq!(false, pt.get_partition(3).unwrap().is_read_only());
        assert_eq!(false, pt.get_partition(4).unwrap().is_read_only());
        assert_eq!(false, pt.get_partition(5).unwrap().is_read_only());

        assert_eq!(false, pt.get_partition(0).unwrap().is_encrypted());
        assert_eq!(false, pt.get_partition(1).unwrap().is_encrypted());
        assert_eq!(false, pt.get_partition(2).unwrap().is_encrypted());
        assert_eq!(false, pt.get_partition(3).unwrap().is_encrypted());
        assert_eq!(false, pt.get_partition(4).unwrap().is_encrypted());
        assert_eq!(false, pt.get_partition(5).unwrap().is_encrypted());
    }

    #[test]
    fn empty_byte_array() {
        let pt = PartitionTable::new(&[]).unwrap();

        assert_eq!(0, pt.len());
        assert!(matches!(pt.get_partition(0), Err(Error::OutOfBounds)));
    }

    #[test]
    fn validation_fails_wo_hash() {
        assert!(matches!(
            PartitionTable::new(&SIMPLE[..RAW_ENTRY_LEN * 3]),
            Err(Error::Invalid)
        ));
    }

    #[test]
    fn validation_fails_wo_hash_max_entries() {
        let mut data = [0u8; PARTITION_TABLE_MAX_LEN];
        for i in 0..96 {
            data[(i * RAW_ENTRY_LEN)..][..RAW_ENTRY_LEN].copy_from_slice(&SIMPLE[..32]);
        }

        assert!(matches!(PartitionTable::new(&data), Err(Error::Invalid)));
    }

    #[test]
    fn validation_succeeds_with_enough_entries() {
        assert_eq!(
            3,
            PartitionTable::new(&SIMPLE[..RAW_ENTRY_LEN * 4])
                .unwrap()
                .len()
        );
    }
}

#[cfg(test)]
mod storage_tests {
    use embedded_storage::{ReadStorage, Storage};

    use super::*;

    struct MockFlash {
        data: [u8; 0x10000],
    }

    impl MockFlash {
        fn new() -> Self {
            let mut data = [23u8; 0x10000];
            data[PARTITION_TABLE_OFFSET as usize..][..PARTITION_TABLE_MAX_LEN as usize]
                .copy_from_slice(include_bytes!("../testdata/single_factory_no_ota.bin"));
            Self { data }
        }
    }

    impl embedded_storage::Storage for MockFlash {
        fn write(&mut self, offset: u32, bytes: &[u8]) -> Result<(), Self::Error> {
            self.data[offset as usize..][..bytes.len()].copy_from_slice(bytes);
            Ok(())
        }
    }

    impl embedded_storage::ReadStorage for MockFlash {
        type Error = crate::partitions::Error;
        fn read(&mut self, offset: u32, buffer: &mut [u8]) -> Result<(), Self::Error> {
            let l = buffer.len();
            buffer[..l].copy_from_slice(&self.data[offset as usize..][..l]);
            Ok(())
        }

        fn capacity(&self) -> usize {
            unimplemented!()
        }
    }

    #[test]
    fn can_read_write_all_of_nvs() {
        let mut storage = MockFlash::new();

        let mut buffer = [0u8; PARTITION_TABLE_MAX_LEN];
        let pt = read_partition_table(&mut storage, &mut buffer).unwrap();

        let nvs = pt
            .find_partition(PartitionType::Data(DataPartitionSubType::Nvs))
            .unwrap()
            .unwrap();
        let mut nvs_partition = nvs.as_embedded_storage(&mut storage);
        assert_eq!(nvs_partition.raw.offset(), 36864);

        assert_eq!(nvs_partition.capacity(), 24576);

        let mut buffer = [0u8; 24576];
        nvs_partition.read(0, &mut buffer).unwrap();
        assert!(buffer.iter().all(|v| *v == 23));
        buffer.fill(42);
        nvs_partition.write(0, &mut buffer).unwrap();
        let mut buffer = [0u8; 24576];
        nvs_partition.read(0, &mut buffer).unwrap();
        assert!(buffer.iter().all(|v| *v == 42));
    }

    #[test]
    fn cannot_read_write_more_than_partition_size() {
        let mut storage = MockFlash::new();

        let mut buffer = [0u8; PARTITION_TABLE_MAX_LEN];
        let pt = read_partition_table(&mut storage, &mut buffer).unwrap();

        let nvs = pt
            .find_partition(PartitionType::Data(DataPartitionSubType::Nvs))
            .unwrap()
            .unwrap();
        let mut nvs_partition = nvs.as_embedded_storage(&mut storage);
        assert_eq!(nvs_partition.raw.offset(), 36864);

        assert_eq!(nvs_partition.capacity(), 24576);

        let mut buffer = [0u8; 24577];
        assert!(nvs_partition.read(0, &mut buffer) == Err(Error::OutOfBounds));
    }
}
