use core::fmt::Write;

use alloc::{string::String, vec::Vec};
use embassy_embedded_hal::{adapter::BlockingAsync, flash::partition::Partition};
use embassy_sync::{blocking_mutex::raw::CriticalSectionRawMutex, mutex::Mutex};
use esp_bootloader_esp_idf::partitions::PartitionTable;
use esp_storage::FlashStorage;
use indoc::writedoc;
use log::info;
use rmk::storage::async_flash_wrapper;
use static_cell::StaticCell;

use crate::PSRAM_ALLOCATOR;

pub struct Partitions {
    pub rmk: Partition<'static, CriticalSectionRawMutex, BlockingAsync<FlashStorage<'static>>>,
    pub acid: Partition<'static, CriticalSectionRawMutex, BlockingAsync<FlashStorage<'static>>>,
}

/// Initialize the flash
pub fn initialize(flash_peripheral: esp_hal::peripherals::FLASH<'static>) -> Partitions {
    static PARTITION_TABLE_BUFFER: StaticCell<Vec<u8, &'static esp_alloc::EspHeap>> =
        StaticCell::new();
    let partition_table_buffer = PARTITION_TABLE_BUFFER.init_with(|| {
        let mut buffer = Vec::<u8, _>::new_in(&PSRAM_ALLOCATOR);
        buffer.resize(
            esp_bootloader_esp_idf::partitions::PARTITION_TABLE_MAX_LEN,
            0_u8,
        );
        buffer
    });

    static FLASH: StaticCell<(
        Mutex<CriticalSectionRawMutex, BlockingAsync<FlashStorage>>,
        PartitionTable<'static>,
    )> = StaticCell::new();
    let (flash, partition_table) = FLASH.init_with(|| {
        let mut flash = FlashStorage::new(flash_peripheral)
            // Flash memory may not be written to while another core is executing from it.
            // By default, `FlashStorage` is configured to abort the operation and log an error message.
            // However, it can also be configured to auto-park the other core, such that writing to
            // flash succeeds.
            // Alternatively, XiP from PSRAM could be used along with the `multicore_ignore` strategy,
            // to avoid having to park the other core, which could result in better performance.
            // Invalid configuration would then present itself as freezing/UB.
            .multicore_auto_park();
        let partition_table = {
            esp_bootloader_esp_idf::partitions::read_partition_table(
                &mut flash,
                partition_table_buffer,
            )
            .expect("Failed to read the partition table.")
        };

        (
            Mutex::<CriticalSectionRawMutex, _>::new(async_flash_wrapper(flash)),
            partition_table,
        )
    });

    {
        let mut buffer = String::new();

        writeln!(buffer, "Partition table:").unwrap();

        for (index, partition) in partition_table.iter().enumerate() {
            writedoc!(
                buffer,
                "
                Partition #{index} {label:?}:
                    offset: 0x{offset:x}
                    length: 0x{len:x}
                    type: 0x{type:?}
                    read only: {read_only}
                    encrypted: {encrypted}
                    magic: {magic}
                ",
                label = partition.label_as_str(),
                offset = partition.offset(),
                len = partition.len(),
                type = partition.partition_type(),
                read_only = partition.is_read_only(),
                encrypted = partition.is_encrypted(),
                magic = partition.magic(),
            )
            .unwrap();
        }

        info!("{}", buffer);
    }
    let flash_part_info_rmk = partition_table
        .iter()
        .find(|partition| partition.label_as_str() == "rmk")
        .expect("No \"rmk\" partition found. Make sure to use the custom partition-table.csv when flashing.");
    let flash_part_info_acid = partition_table
        .iter()
        .find(|partition| partition.label_as_str() == "acid")
        .expect("No \"acid\" partition found. Make sure to use the custom partition-table.csv when flashing.");

    Partitions {
        rmk: Partition::new(
            flash,
            flash_part_info_rmk.offset(),
            flash_part_info_rmk.len(),
        ),
        acid: Partition::new(
            flash,
            flash_part_info_acid.offset(),
            flash_part_info_acid.len(),
        ),
    }
}