// src/flash.rs
//! Flash memory module (STM32U575, minimal non-secure subset)
//!
//! Included:
//! - Constrain into parts (register proxies)
//! - Unlock/lock (non-secure) via NSKEYR / NSCR
//! - Page erase using NSCR.PER/PNB/BKER/START
//! - 128-bit (quad-word) program using NSCR.PG and 4x32 writes
//! - Simple wait (polling on BSY), minimal EOP clear
//! - Byte-granular write built on top of 16-byte chunks
//!
//! Skipped / TODO:
//! - Secure programming path (SECKEYR/SECCR/SECSR)
//! - Burst program (8x quad-word)
//! - Detailed error decoding (PROGERR/PGSERR/WRPERR...)
//! - Option bytes, protections, ECC interactions
//! - Exact geometry detection: we assume 8 KiB pages, 256 pages total

#![deny(missing_docs)]

use crate::pac as pac;
use crate::pac::FLASH;
use core::{mem, ptr};

/// Simple Flash error set for this minimal port
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum Error {
    /// Flash controller is busy
    Busy,
    /// Operation failed (generic)
    Failure,
    /// Address alignment issue
    Unaligned,
    /// Page index out of range
    PageOutOfRange,
}

/// A Flash page index wrapper (0-based)
///
/// For STM32U575 this example assumes 8 KiB pages, 256 total pages
/// (128 per bank). Adjust if your device differs.
#[derive(Copy, Clone, Debug)]
pub struct FlashPage(pub u16);

impl FlashPage {
    /// Convert page index to absolute address
    ///
    /// 0x0800_0000 + page_idx * PAGE_SIZE
    pub const fn to_address(&self) -> usize {
        FLASH_BASE as usize + (self.0 as usize) * PAGE_SIZE
    }
}

/// Base of Flash memory
pub const FLASH_BASE: u32 = 0x0800_0000;
/// Page size in bytes (assumed)
pub const PAGE_SIZE: usize = 8 * 1024; // 8 KiB
/// Total pages assumed
pub const TOTAL_PAGES: u16 = 256;
/// Pages per bank assumed
pub const PAGES_PER_BANK: u16 = TOTAL_PAGES / 2;

/// Extension trait to constrain the FLASH peripheral
pub trait FlashExt {
    /// Constrain the FLASH peripheral to play nicely with the other abstractions
    fn constrain(self) -> Parts;
}

impl FlashExt for FLASH {
    fn constrain(self) -> Parts {
        Parts {
            acr: ACR,
            nskeyr: NSKEYR,
            optkeyr: OPTKEYR,
            nssr: NSSR,
            nscr: NSCR,
            eccr: ECCR,
            pdkey1r: PDKEY1R,
            pdkey2r: PDKEY2R,
        }
    }
}

/// Constrained FLASH peripheral
pub struct Parts {
    /// Opaque ACR register
    pub acr: ACR,
    /// Opaque NSKEYR register (non-secure key)
    pub nskeyr: NSKEYR,
    /// Opaque OPTKEYR register
    pub optkeyr: OPTKEYR,
    /// Opaque NSSR register (non-secure status)
    pub nssr: NSSR,
    /// Opaque NSCR register (non-secure control)
    pub nscr: NSCR,
    /// Opaque ECCR register
    pub eccr: ECCR,
    /// Opaque PDKEY1R register
    pub pdkey1r: PDKEY1R,
    /// Opaque PDKEY2R register
    pub pdkey2r: PDKEY2R,
}

macro_rules! generate_register {
    ($A:ident, $field:ident, $name:expr) => {
        #[doc = "Opaque "]
        #[doc = $name]
        #[doc = " register"]
        pub struct $A;

        impl $A {
            #[allow(unused)]
            #[inline(always)]
            pub(crate) fn reg(&mut self) -> &pac::flash::$A {
                // NOTE(unsafe): this proxy grants exclusive access to this register
                unsafe { &(*FLASH::ptr()).$field() }
            }
        }
    };
    ($A:ident, $field:ident) => {
        generate_register!($A, $field, stringify!($A));
    };
}

// Register proxies (names match stm32u575 PAC)
generate_register!(ACR, acr);
generate_register!(NSKEYR, nskeyr);
generate_register!(OPTKEYR, optkeyr);
generate_register!(NSSR, nssr);
generate_register!(NSCR, nscr);
generate_register!(ECCR, eccr);
generate_register!(PDKEY1R, pdkey1r);
generate_register!(PDKEY2R, pdkey2r);

// Unlock keys (same values as STM32U5 HAL)
const FLASH_KEY1: u32 = 0x4567_0123;
const FLASH_KEY2: u32 = 0xCDEF_89AB;

/// Active programming session on FLASH (non-secure)
///
/// Drop locks the controller automatically.
pub struct FlashProgramming<'a> {
    nssr: &'a mut NSSR,
    nscr: &'a mut NSCR,
}

impl<'a> Drop for FlashProgramming<'a> {
    fn drop(&mut self) {
        // Lock on drop
        self.lock();
    }
}

impl Parts {
    /// Unlock the flash non-secure control registers via NSKEYR for programming/erase
    pub fn unlock(&mut self) -> Result<FlashProgramming<'_>, Error> {
        // Write keys to NSKEYR
        unsafe {
            self.nskeyr.reg().write(|w| w.bits(FLASH_KEY1));
            self.nskeyr.reg().write(|w| w.bits(FLASH_KEY2));
        }

        // Check NSCR.LOCK bit
        let locked = self.nscr.reg().read().lock().bit_is_set();
        if locked {
            Err(Error::Failure)
        } else {
            Ok(FlashProgramming {
                nssr: &mut self.nssr,
                nscr: &mut self.nscr,
            })
        }
    }
}

impl<'a> FlashProgramming<'a> {
    /// Lock the flash memory controller (NSCR.LOCK=1)
    pub fn lock(&mut self) {
        self.nscr.reg().modify(|_, w| w.lock().set_bit());
    }

    /// Poll for completion (busy clear).
    ///
    /// Minimal: we only look at BSY. If BSY is clear, we return Ok.
    /// You can extend this to decode error flags and EOP as needed.
    fn wait(&mut self) -> Result<(), Error> {
        while self.nssr.reg().read().bsy().bit_is_set() {}
        Ok(())
    }

    /// Clear EOP if set (optional; helps keep SR clean between operations).
    fn clear_eop(&mut self) {
        let sr = self.nssr.reg().read();
        if sr.eop().bit_is_set() {
            // Write-one-to-clear EOP
            unsafe {
                self.nssr.reg().write(|w| w.eop().set_bit());
            }
        }
    }

    /// Erase a single page by index.
    ///
    /// Assumes:
    /// - TOTAL_PAGES = 256
    /// - PAGES_PER_BANK = 128
    /// - BKER=0 => bank1, BKER=1 => bank2
    pub fn erase_page(&mut self, page: FlashPage) -> Result<(), Error> {
        let idx = page.0;
        if idx >= TOTAL_PAGES {
            return Err(Error::PageOutOfRange);
        }

        let (bker_set, pnb) = if idx < PAGES_PER_BANK {
            (false, idx as u8)
        } else {
            (true, (idx - PAGES_PER_BANK) as u8)
        };

        // PER=1, set PNB and BKER, then START=1
        unsafe {
            self.nscr.reg().modify(|_, w| {
                w.per().set_bit();
                if bker_set {
                    w.bker().set_bit();
                } else {
                    w.bker().clear_bit();
                }
                w.pnb().bits(pnb);
                w
            });
            self.nscr.reg().modify(|_, w| w.strt().set_bit());
        }

        self.wait()?;
        self.clear_eop();

        // Clear PER
        self.nscr.reg().modify(|_, w| w.per().clear_bit());

        Ok(())
    }

    /// Program one 128-bit (quad-word) at the given address.
    ///
    /// - Address must be 16-byte aligned.
    /// - Data are provided as 16 raw bytes; we program four u32 words.
    pub fn write_quadword(&mut self, address: usize, data16: &[u8; 16]) -> Result<(), Error> {
        if (address & 0xF) != 0 {
            return Err(Error::Unaligned);
        }

        // Extract four 32-bit words (little-endian)
        let w0 = u32::from_le_bytes([data16[0], data16[1], data16[2], data16[3]]);
        let w1 = u32::from_le_bytes([data16[4], data16[5], data16[6], data16[7]]);
        let w2 = u32::from_le_bytes([data16[8], data16[9], data16[10], data16[11]]);
        let w3 = u32::from_le_bytes([data16[12], data16[13], data16[14], data16[15]]);

        // PG=1
        self.nscr.reg().modify(|_, w| w.pg().set_bit());

        // Perform 4x 32-bit writes
        let mut dst = address as *mut u32;
        unsafe {
            ptr::write_volatile(dst, w0);
            dst = dst.add(1);
            ptr::write_volatile(dst, w1);
            dst = dst.add(1);
            ptr::write_volatile(dst, w2);
            dst = dst.add(1);
            ptr::write_volatile(dst, w3);
        }

        self.wait()?;
        self.clear_eop();

        // PG=0
        self.nscr.reg().modify(|_, w| w.pg().clear_bit());

        Ok(())
    }

    /// Write arbitrary bytes, rounding to 16-byte chunks.
    ///
    /// For leading/trailing partial chunks, we fill missing bytes with 0xFF,
    /// which matches erased Flash default.
    pub fn write(&mut self, address: usize, buf: &[u8]) -> Result<(), Error> {
        // Handle leading unaligned portion
        let align_off = address & 0xF;
        let mut cur_addr = address;

        if align_off != 0 {
            let head = core::cmp::min(16 - align_off, buf.len());
            let mut q = [0xFFu8; 16];
            // Pull existing erased fill (0xFF), then overlay provided bytes
            q[align_off..align_off + head].copy_from_slice(&buf[..head]);
            self.write_quadword(address - align_off, &q)?;
            cur_addr += head;
        }

        // Aligned middle
        let aligned_len = if cur_addr > address {
            buf.len() - (cur_addr - address)
        } else {
            buf.len()
        };
        let aligned_start = buf.len() - aligned_len + (cur_addr - address);
        let mut i = 0;
        while i + 16 <= aligned_len {
            let mut q = [0u8; 16];
            q.copy_from_slice(
                &buf[aligned_start + i..aligned_start + i + 16],
            );
            self.write_quadword(cur_addr, &q)?;
            cur_addr += 16;
            i += 16;
        }

        // Trailing remainder
        let rem = aligned_len - i;
        if rem > 0 {
            let mut q = [0xFFu8; 16];
            q[..rem].copy_from_slice(
                &buf[aligned_start + i..aligned_start + i + rem],
            );
            self.write_quadword(cur_addr, &q)?;
        }

        Ok(())
    }

    /// Mass erase (both banks). Minimal example; adjust to your needs.
    pub fn erase_all(&mut self) -> Result<(), Error> {
        // Set MER1 and MER2, then START
        self.nscr.reg().modify(|_, w| w.mer1().set_bit());
        self.nscr.reg().modify(|_, w| w.mer2().set_bit());
        self.nscr.reg().modify(|_, w| w.strt().set_bit());

        self.wait()?;
        self.clear_eop();

        // Clear MER bits
        self.nscr.reg().modify(|_, w| w.mer1().clear_bit());
        self.nscr.reg().modify(|_, w| w.mer2().clear_bit());

        Ok(())
    }
}
