//! Flash memory, with major reference from STM32F7xx_HAL crate,
//! as well as @Astro 's work in STM32F4xx_HAL

use stm32h7xx_hal::pac::FLASH;

/// Base address of flash memory on AXIM interface.
const FLASH_BASE: *mut u8 = 0x800_0000 as *mut u8;

/// The last valid flash address in STM32H743
const MAX_FLASH_ADDRESS: *mut u8 = 0x81F_FFFF as *mut u8;

/// Flash programming error.
#[derive(Debug, PartialEq, Eq)]
pub enum Error {
    Busy,
    Locked,
    WriteProtection,
    ProgrammingSequence,
    Strobe,
    Inconsistency,
    Operation,
    ErrorCorrectionCode,
    ReadProtection,
    ReadSecure,
    WriteError,
}

/// Embedded flash memory.
pub struct Flash {
    registers: FLASH,
}

impl Flash {
    /// Creates a new Flash instance.
    pub fn new(flash: FLASH) -> Self {
        let mut flash = Self { registers: flash };
        // Lock both banks initially
        flash.lock();
        flash
    }

    /// Unlocks the FLASH_CR1/2 register.
    pub fn unlock(&mut self) {
        // Note: Unlocking an unlocked bank will cause HardFault

        // Unlock bank 1 if needed.
        if self.bank1_is_locked() {
            self.registers
                .bank1_mut()
                .keyr
                .write(|w| unsafe { w.keyr().bits(0x45670123) });
            self.registers
                .bank1_mut()
                .keyr
                .write(|w| unsafe { w.keyr().bits(0xCDEF89AB) });
        }

        // Unlock bank 2 if needed.
        if self.bank2_is_locked() {
            self.registers
                .bank2_mut()
                .keyr
                .write(|w| unsafe { w.keyr().bits(0x45670123) });
            self.registers
                .bank2_mut()
                .keyr
                .write(|w| unsafe { w.keyr().bits(0xCDEF89AB) });
        }
    }

    /// Unlocks the FLASH_OPTCR register
    pub fn unlock_optcr(&mut self) {
        if self.optcr_is_locked() {
            self.registers
                .optkeyr_mut()
                .write(|w| unsafe { w.optkeyr().bits(0x08192A3B) });
            self.registers
                .optkeyr_mut()
                .write(|w| unsafe { w.optkeyr().bits(0x4C5D6E7F) });
        }
    }

    /// Locks the FLASH_CR1/2 register.
    pub fn lock(&mut self) {
        self.registers
            .bank1_mut()
            .cr
            .modify(|_, w| w.lock().set_bit());
        self.registers
            .bank2_mut()
            .cr
            .modify(|_, w| w.lock().set_bit());
    }

    /// Lock the FLASH_OPTCR register
    pub fn lock_optcr(&mut self) {
        self.registers
            .optcr_mut()
            .modify(|_, w| w.optlock().set_bit());
    }

    // More literal methods to get bank status
    fn bank1_is_locked(&self) -> bool {
        self.registers.bank1_mut().cr.read().lock().bit_is_set()
    }
    fn bank2_is_locked(&self) -> bool {
        self.registers.bank2_mut().cr.read().lock().bit_is_set()
    }
    fn optcr_is_locked(&self) -> bool {
        self.registers.optcr().read().optlock().bit_is_set()
    }

    pub fn is_locked(&self) -> bool {
        self.bank1_is_locked() || self.bank2_is_locked()
    }

    /// Returns `true` if a write/erase operation is in progress.
    fn is_busy(&self) -> bool {
        let (sr1, sr2) = (
            self.registers.bank1_mut().sr.read(),
            self.registers.bank2_mut().sr.read(),
        );
        sr1.bsy().bit_is_set() || sr2.bsy().bit_is_set()
    }

    /// Returns `true` if a write/erase operation is in a command queue buffer.
    fn is_queuing(&self) -> bool {
        let (sr1, sr2) = (
            self.registers.bank1_mut().sr.read(),
            self.registers.bank2_mut().sr.read(),
        );
        sr1.qw().bit_is_set() || sr2.qw().bit_is_set()
    }

    /// Returns `true` if write buffer is not empty in bank 1
    fn bank1_is_buffering(&self) -> bool {
        let sr1 = self.registers.bank1_mut().sr.read();
        sr1.wbne().bit_is_set()
    }

    /// Returns `true` if write buffer is not empty in bank 2
    fn bank2_is_buffering(&self) -> bool {
        let sr2 = self.registers.bank2_mut().sr.read();
        sr2.wbne().bit_is_set()
    }

    /// Erase a sector.
    pub fn erase_sector(&mut self, bank_number: u8, sector_number: u8) -> Result<(), Error> {
        // Assert parameters validity
        assert!(bank_number == 1 || bank_number == 2);
        assert!(sector_number < 8);
        // 1. Check & clear error flags
        self.clear_errors();
        // 2. Unlock FLASH_CR1/2 register if necessary
        self.unlock();
        // 3. Set SER1/2 & SNB1/2 bits in FLASH_CR1/2 register
        // 4. Set START1/2 bit in FLASH_CR1/2 register
        match bank_number {
            1 => {
                self.registers
                    .bank1_mut()
                    .cr
                    .modify(|_, w| unsafe { w.ser().set_bit().snb().bits(sector_number) });
                self.registers
                    .bank1_mut()
                    .cr
                    .modify(|_, w| w.start().set_bit());
            }
            2 => {
                self.registers
                    .bank2_mut()
                    .cr
                    .modify(|_, w| unsafe { w.ser().set_bit().snb().bits(sector_number) });
                self.registers
                    .bank2_mut()
                    .cr
                    .modify(|_, w| w.start().set_bit());
            }
            _ => unreachable!(),
        }
        // Lock the flash CR again
        self.lock();
        Ok(())
    }

    /// Erase a bank.
    pub fn erase_bank(&mut self, bank_number: u8) -> Result<(), Error> {
        // Assert parameters validity
        assert!(bank_number == 1 || bank_number == 2);
        // 1. Check & clear error flags
        self.clear_errors();
        // 2. Unlock FLASH_CR1/2 register if necessary
        self.unlock();
        // 3 & 4. Set BER1/2 bit and START1/2 bit in FLASH_CR1/2 register
        //         Wait until the corresponding QW1/2 bit is cleared.
        // Note: By setting BER1/2 bit alongside START1/2,
        //       mass erase is invoked since it supersedes sector erase.
        match bank_number {
            1 => {
                self.registers
                    .bank1_mut()
                    .cr
                    .modify(|_, w| w.ber().set_bit().start().set_bit());
                while self.registers.bank1_mut().sr.read().qw().bit_is_set() {}
            }
            2 => {
                self.registers
                    .bank2_mut()
                    .cr
                    .modify(|_, w| w.ber().set_bit().start().set_bit());
                while self.registers.bank2_mut().sr.read().qw().bit_is_set() {}
            }
            _ => unreachable!(),
        }
        // Lock the flash CR again
        self.lock();
        Ok(())
    }

    /// Mass erases of the flash memory.
    pub fn mass_erase(&mut self) -> Result<(), Error> {
        // 1. Check & clear error flags
        self.clear_errors();
        // 2. Unlock Flash_CR1&2, FLASH_OPTCR registers
        self.unlock();
        self.unlock_optcr();
        // 3. Set MER in FLASH_OPTCR to 1, wait until both QW to clear
        self.registers.optcr_mut().modify(|_, w| w.mer().set_bit());
        while self.is_queuing() {}
        // Lock the flash CR and OPTCR again
        self.lock();
        self.lock_optcr();
        Ok(())
    }

    /// Program flash words (32-bytes).
    /// Flashing incomplete flash word is "tolerated", but you have been warned..
    pub fn program<'a, 'b>(&'a mut self, start_offset: usize, data: &'b [u8]) -> Result<(), Error> {
        if (start_offset % 32 != 0) || (data.len() % 32 != 0) {
            log::warn!("Warning: This flash operation might not be supported...");
            log::warn!("Consider force writing the data in buffer...");
        }
        self.clear_errors();
        // Invoke single writes per 32-bytes row
        let mut current_address = start_offset;
        let mut remaining_data = data;
        while remaining_data.len() != 0 {
            let single_write_size = 32 - (current_address % 32);
            // Determine the index that split the coming row and the remaining bytes.
            let splitting_index = core::cmp::min(single_write_size, remaining_data.len());
            let single_row_data = &remaining_data[..splitting_index];
            // 1. Unlock FLASH_CR1/2 register if necessary
            self.unlock();
            // 2. Set PG bit in FLASH_CR1/2
            self.registers
                .bank1_mut()
                .cr
                .modify(|_, w| w.pg().set_bit());
            self.registers
                .bank2_mut()
                .cr
                .modify(|_, w| w.pg().set_bit());
            // 3. Check Protection
            // There should not be any data protection anyway...
            // 4. Write data byte by byte
            for (index, byte) in single_row_data.iter().enumerate() {
                while self.is_busy() {}
                match self.check_errors() {
                    Ok(_) => {
                        let address: *mut u8 = unsafe { FLASH_BASE.add(current_address + index) };
                        if address > MAX_FLASH_ADDRESS {
                            self.registers
                                .bank1_mut()
                                .cr
                                .modify(|_, w| w.pg().clear_bit());
                            self.registers
                                .bank2_mut()
                                .cr
                                .modify(|_, w| w.pg().clear_bit());
                            return Err(Error::WriteError);
                        } else {
                            unsafe {
                                core::ptr::write_volatile(address, *byte);
                            }
                        }
                    }
                    Err(error) => {
                        self.registers
                            .bank1_mut()
                            .cr
                            .modify(|_, w| w.pg().clear_bit());
                        self.registers
                            .bank2_mut()
                            .cr
                            .modify(|_, w| w.pg().clear_bit());
                        return Err(error);
                    }
                }
            }
            // 5. Wait till the command queue of the device to be empty
            while self.is_queuing() {}
            // Modify remaining data and the address for next 32-bytes row.
            remaining_data = &remaining_data[splitting_index..];
            current_address += single_row_data.len();
        }
        // Reset PG1/2
        self.registers
            .bank1_mut()
            .cr
            .modify(|_, w| w.pg().clear_bit());
        self.registers
            .bank2_mut()
            .cr
            .modify(|_, w| w.pg().clear_bit());
        // Lock FLASH_CR1/2 register
        self.lock();
        Ok(())
    }

    /// Force empty the bytes buffer for flash programming
    /// Warning: It can invalidate the whole flash due to invalid CRC.
    pub fn force_write<'a>(&'a mut self) -> Result<(), Error> {
        if self.bank1_is_buffering() {
            self.registers
                .bank1_mut()
                .cr
                .modify(|_, w| w.fw().set_bit());
        }
        if self.bank2_is_buffering() {
            self.registers
                .bank2_mut()
                .cr
                .modify(|_, w| w.fw().set_bit());
        }
        Ok(())
    }

    /// Read a slice from flash memory
    pub fn read(&self, start_offset: usize, len: usize) -> &'static [u8] {
        let address = unsafe { FLASH_BASE.add(start_offset) };
        unsafe { core::slice::from_raw_parts(address, len) }
    }

    /// Releases the flash peripheral.
    pub fn free(self) -> FLASH {
        self.registers
    }

    /// Checks the error flags.
    fn check_errors(&self) -> Result<(), Error> {
        let (sr1, sr2) = (
            self.registers.bank1_mut().sr.read(),
            self.registers.bank2_mut().sr.read(),
        );

        if sr1.wrperr().bit_is_set() || sr2.wrperr().bit_is_set() {
            Err(Error::WriteProtection)
        } else if sr1.pgserr().bit_is_set() || sr2.pgserr().bit_is_set() {
            Err(Error::ProgrammingSequence)
        } else if sr1.strberr().bit_is_set() || sr2.strberr().bit_is_set() {
            Err(Error::Strobe)
        } else if sr1.incerr().bit_is_set() || sr2.incerr().bit_is_set() {
            Err(Error::Inconsistency)
        } else if sr1.operr().bit_is_set() || sr2.operr().bit_is_set() {
            Err(Error::Operation)
        } else if sr1.sneccerr1().bit_is_set()
            || sr1.dbeccerr().bit_is_set()
            || sr2.sneccerr1().bit_is_set()
            || sr2.dbeccerr().bit_is_set()
        {
            Err(Error::ErrorCorrectionCode)
        } else if sr1.rdperr().bit_is_set() || sr2.rdperr().bit_is_set() {
            Err(Error::ReadProtection)
        } else if sr1.rdserr().bit_is_set() || sr2.rdserr().bit_is_set() {
            Err(Error::ReadSecure)
        } else {
            Ok(())
        }
    }

    /// Clears all error flags.
    fn clear_errors(&mut self) {
        self.registers.bank1_mut().ccr.write(|w| {
            w.clr_wrperr()
                .set_bit()
                .clr_pgserr()
                .set_bit()
                .clr_strberr()
                .set_bit()
                .clr_incerr()
                .set_bit()
                .clr_operr()
                .set_bit()
                .clr_rdperr()
                .set_bit()
                .clr_rdserr()
                .set_bit()
                .clr_sneccerr()
                .set_bit()
                .clr_dbeccerr()
                .set_bit()
        });
        self.registers.bank2_mut().ccr.write(|w| {
            w.clr_wrperr()
                .set_bit()
                .clr_pgserr()
                .set_bit()
                .clr_strberr()
                .set_bit()
                .clr_incerr()
                .set_bit()
                .clr_operr()
                .set_bit()
                .clr_rdperr()
                .set_bit()
                .clr_rdserr()
                .set_bit()
                .clr_sneccerr()
                .set_bit()
                .clr_dbeccerr()
                .set_bit()
        });
    }
}