#![no_std]

use imxrt_drivers_flexspi as flexspi;
pub use imxrt_drivers_flexspi::{SeqId, SeqIdRange, Sequence};
use ral_registers as ral;

pub mod sequences {
    pub mod common;
}

/// Helpers for interacting with flash.
pub mod flash {
    use super::{IpCmd, SeqId, flexspi};

    pub mod adesto;
    pub mod issi;
    pub mod winbond;

    const READ: IpCmd = IpCmd::new(SeqId::Seq00, 1).unwrap();
    const READ_STATUS: IpCmd = IpCmd::new(SeqId::Seq01, 1).unwrap();
    const WRITE_ENABLE: IpCmd = IpCmd::new(SeqId::Seq03, 1).unwrap();
    const ERASE_SECTOR: IpCmd = IpCmd::new(SeqId::Seq05, 2).unwrap();
    const PAGE_PROGRAM: IpCmd = IpCmd::new(SeqId::Seq09, 2).unwrap();
    const CHIP_ERASE: IpCmd = IpCmd::new(SeqId::Seq11, 2).unwrap();

    const SET_READ_PARAMS: IpCmd = IpCmd::new(SeqId::Seq02, 1).unwrap();
    const RESET: IpCmd = IpCmd::new(SeqId::Seq07, 2).unwrap();

    const PAGE_SIZE_BYTES: usize = 256;

    const STATUS_WIP: u8 = 1 << 0;

    /// Write data to the flash array, starting at `flash_start`.
    ///
    /// The implementation handles page crossings and writes into the
    /// middle of pages. The call blocks until the peripheral and flash
    /// array is idle.
    ///
    /// This call does not eagerly erase any flash sectors. You'll need
    /// to do that yourself.
    pub fn write(flexspi: flexspi::Instance, mut flash_start: usize, data: &[u8]) {
        for page in aligned_chunks(flash_start, data, PAGE_SIZE_BYTES) {
            crate::start_ip_cmd(flexspi, PAGE_PROGRAM, flash_start, page);
            crate::transmit_bytes(flexspi, page);
            crate::wait_for_ip_cmd_done(flexspi);
            crate::clear_tx_fifo(flexspi);
            crate::wait_for_idle(flexspi);

            flash_start = flash_start.saturating_add(page.len());

            wait_for_wip_clear(flexspi);
        }
    }

    /// Poll the status register while WIP is set.
    fn wait_for_wip_clear(flexspi: flexspi::Instance) {
        while {
            let status = read_status(flexspi);
            status & STATUS_WIP != 0
        } {}
    }

    /// Read the status register.
    fn read_status(flexspi: flexspi::Instance) -> u8 {
        let mut status = [0_u8; 1];
        crate::start_ip_cmd(flexspi, READ_STATUS, 0, &status);
        crate::receive_bytes(flexspi, &mut status);
        crate::wait_for_ip_cmd_done(flexspi);
        crate::clear_rx_fifo(flexspi);
        crate::wait_for_idle(flexspi);
        status[0]
    }

    /// Read data from the flash array starting at `flash_start`.
    ///
    /// The implementation handles reads into the middle of pages, along
    /// with page crossings. The call blocks until all bytes are in the
    /// `data` buffer.
    pub fn read(flexspi: flexspi::Instance, flash_start: usize, data: &mut [u8]) {
        crate::start_ip_cmd(flexspi, READ, flash_start, data);
        crate::receive_bytes(flexspi, data);
        crate::wait_for_ip_cmd_done(flexspi);
        crate::clear_rx_fifo(flexspi);
        crate::wait_for_idle(flexspi);
    }

    /// Erase a sector.
    ///
    /// `flash_start` is an address in the flash part. The call
    /// erases the entire sector, even if `flash_start` isn't a
    /// sector aligned address. Blocks until the flash array is
    /// idle.
    pub fn erase_sector(flexspi: flexspi::Instance, flash_start: usize) {
        crate::start_ip_cmd(flexspi, ERASE_SECTOR, flash_start, &[]);
        crate::wait_for_ip_cmd_done(flexspi);
        crate::wait_for_idle(flexspi);

        wait_for_wip_clear(flexspi);
    }

    /// Erase the entire chip.
    pub fn erase_chip(flexspi: flexspi::Instance) {
        crate::start_ip_cmd(flexspi, CHIP_ERASE, 0, &[]);
        crate::wait_for_ip_cmd_done(flexspi);
        crate::wait_for_idle(flexspi);

        wait_for_wip_clear(flexspi);
    }

    /// Reset the flash chip.
    pub fn reset(flexspi: flexspi::Instance) {
        crate::start_ip_cmd(flexspi, RESET, 0, &[]);
        crate::wait_for_ip_cmd_done(flexspi);
        crate::wait_for_idle(flexspi);

        wait_for_wip_clear(flexspi);
    }

    /// Produce chunks of bytes suitable for page aligned writing.
    fn aligned_chunks(start: usize, bytes: &[u8], page_size: usize) -> impl Iterator<Item = &[u8]> {
        let next_page_start = page_size - (start % page_size);
        core::iter::once(&bytes[..next_page_start])
            .chain(bytes[next_page_start..].chunks(page_size))
    }

    #[cfg(test)]
    mod tests {
        extern crate std;
        use std::vec::Vec;

        use super::aligned_chunks;

        fn make_buffer() -> Vec<u8> {
            let mut buffer = Vec::with_capacity(1024);
            for i in 1..=4 {
                buffer.extend(std::iter::repeat(i).take(256));
            }
            assert_eq!(buffer.len(), 1024);
            buffer
        }

        #[test]
        fn aligned() {
            let buffer = make_buffer();
            let mut chunks = aligned_chunks(256, &buffer, 256);
            assert_eq!(chunks.next().unwrap(), &[1; 256][..]);
            assert_eq!(chunks.next().unwrap(), &[2; 256][..]);
            assert_eq!(chunks.next().unwrap(), &[3; 256][..]);
            assert_eq!(chunks.next().unwrap(), &[4; 256][..]);
            assert_eq!(chunks.next(), None);
        }

        #[test]
        fn unaligned_by_one() {
            let buffer = make_buffer();
            let mut chunks = aligned_chunks(1, &buffer, 256);
            assert_eq!(chunks.next().unwrap(), &buffer[0..255]);
            assert_eq!(chunks.next().unwrap(), &buffer[255..511]);
            assert_eq!(chunks.next().unwrap(), &buffer[511..767]);
            assert_eq!(chunks.next().unwrap(), &buffer[767..1023]);
            assert_eq!(chunks.next().unwrap(), &buffer[1023..]);
            assert_eq!(chunks.next(), None);
        }

        #[test]
        fn unalged_by_almost_a_page() {
            let buffer = make_buffer();
            let mut chunks = aligned_chunks(255, &buffer, 256);
            assert_eq!(chunks.next().unwrap(), &buffer[0..1]);
            assert_eq!(chunks.next().unwrap(), &buffer[1..257]);
            assert_eq!(chunks.next().unwrap(), &buffer[257..513]);
            assert_eq!(chunks.next().unwrap(), &buffer[513..769]);
            assert_eq!(chunks.next().unwrap(), &buffer[769..1024]);
            assert_eq!(chunks.next(), None);
        }
    }
}

pub mod imxrt10xx;
pub mod imxrt11xx;

/// A handle for an IP command sequence.
#[derive(Clone, Copy, PartialEq, Eq)]
pub struct IpCmd {
    seq_id: SeqId,
    seq_num: usize,
}

impl defmt::Format for IpCmd {
    fn format(&self, fmt: defmt::Formatter) {
        defmt::write!(
            fmt,
            "IpCmd({=usize},{=usize})",
            self.seq_id as usize,
            self.seq_num
        )
    }
}

impl IpCmd {
    /// Create an IP command handle for a command sequence.
    pub const fn for_sequences(seq_id: SeqId, seqs: &[Sequence]) -> Option<Self> {
        Self::new(seq_id, seqs.len())
    }

    /// Create an IP command handle for executing the given
    /// sequence.
    pub const fn new(seq_id: SeqId, seq_num: usize) -> Option<Self> {
        if seq_num == 0
            || seq_num > flexspi::MAX_SEQ_PER_IP_CMD
            || (seq_id as usize) + seq_num > flexspi::LUT_SIZE
        {
            None
        } else {
            Some(IpCmd { seq_id, seq_num })
        }
    }
}

/// Reset the peripheral.
///
/// Specify how large your flash part is, in KiB. Also
/// specify how large your TX and RX FIFOs are for your
/// FlexSPI peripheral.
///
/// After this call, you should install IP command sequences.
pub fn reset(flexspi: flexspi::Instance, flash_size_kib: usize, fifo_capacity_bytes: usize) {
    ral::write_reg!(flexspi, flexspi, MCR0, SWRESET: 1);
    while ral::read_reg!(flexspi, flexspi, MCR0, SWRESET == 1) {}

    ral::modify_reg!(flexspi, flexspi, MCR0, MDIS: 1);

    ral::modify_reg!(flexspi, flexspi, MCR0,
        AHBGRANTWAIT: 0xFF,
        IPGRANTWAIT: 0xFF,
        SCKFREERUNEN: 0,
        COMBINATIONEN: 0,
        DOZEEN: 0,
        HSEN: 0,
        SERCLKDIV: DIVIDE_1,
        ATDFEN: IP_BUS,
        ARDFEN: IP_BUS,
        RXCLKSRC: LOOPBACK_DQS,
    );
    ral::write_reg!(flexspi, flexspi, MCR1, SEQWAIT: 0xFFFF, AHBBUSWAIT: 0xFFFF);
    ral::write_reg!(flexspi, flexspi, INTEN, 0);
    ral::write_reg!(flexspi, flexspi, INTR, u32::MAX);
    ral::write_reg!(
        flexspi,
        flexspi,
        FLSHCR0[flexspi::A1],
        flash_size_kib as u32
    );

    ral::write_reg!(flexspi, flexspi, IPRXFCR, RXWMRK: 0, RXDMAEN: 0, CLRIPRXF: 1);

    let tfdr_bytes = tfdr_len(flexspi) * size_of::<u32>();
    let txmrk_bytes = tfdr_bytes.min(fifo_capacity_bytes);
    let txwmrk = txmrk_bytes.div_ceil(size_of::<u64>()).saturating_sub(1);
    ral::write_reg!(flexspi, flexspi, IPTXFCR, TXWMRK: txwmrk as u32, TXDMAEN: 0, CLRIPTXF: 1);

    ral::modify_reg!(flexspi, flexspi, MCR0, MDIS: 0);
}

/// Install the a command sequence for an IP command.
///
/// Assumes that the peripheral is idle. You should install
/// sequences after reset.
///
/// Returns the handle to the command sequence. Note that you
/// can pre-compute this sequence and check it against the
/// return. However, the handle isn't valid until this call.
pub fn install_ip_cmd(
    flexspi: flexspi::Instance,
    seq_id: SeqId,
    sequences: &[Sequence],
) -> Option<IpCmd> {
    let ip_cmd = IpCmd::for_sequences(seq_id, sequences)?;

    flexspi::unlock_lut(flexspi);
    for (seq_id, seq) in SeqIdRange::start(seq_id).zip(sequences) {
        flexspi::set_sequence(flexspi, seq_id, seq);
    }
    flexspi::lock_lut(flexspi);

    Some(ip_cmd)
}

/// Start a previously-installed IP command.
///
/// Generally, this call tries to minimize how long it blocks the CPU.
/// If you're transmitting or receiving data, you should perform that
/// operation as soon as this call returns.
///
/// If you're reading / writing from a register, set flash start to
/// zero. If you have nothing to read or write, set flash start to
/// zero and supply an empty buffer.
///
/// Call assumes that there is no in-progress IP command.
pub fn start_ip_cmd(flexspi: flexspi::Instance, ip_cmd: IpCmd, flash_start: usize, data: &[u8]) {
    ral::write_reg!(flexspi, flexspi, IPCR0, flash_start as u32);
    ral::write_reg!(flexspi, flexspi, IPCR1,
        IDATSZ: data.len() as u32,
        ISEQNUM: ip_cmd.seq_num.saturating_sub(1) as u32,
        ISEQID: ip_cmd.seq_id as u32,
    );
    ral::write_reg!(flexspi, flexspi, IPCMD, TRG: 1);
}

/// Clear the TX FIFO.
#[inline]
pub fn clear_tx_fifo(flexspi: flexspi::Instance) {
    ral::modify_reg!(flexspi, flexspi, IPTXFCR, CLRIPTXF: 1);
}

/// Clear the RX FIFO.
#[inline]
pub fn clear_rx_fifo(flexspi: flexspi::Instance) {
    ral::modify_reg!(flexspi, flexspi, IPRXFCR, CLRIPRXF: 1);
}

/// Transmit bytes to the flash part.
///
/// Returns as soon as all data has reached the transmit FIFO.
/// You must then wait for the IP command to complete then clear
/// the transmit FIFO.
pub fn transmit_bytes(flexspi: flexspi::Instance, data: &[u8]) {
    // Reset clamps the watermark to the maximum size of the
    // TFDR array.
    let txwmrk = ral::read_reg!(flexspi, flexspi, IPTXFCR, TXWMRK).saturating_add(1) as usize;
    let len = txwmrk * const { size_of::<u64>() / size_of::<u32>() };

    let words = data.chunks(size_of::<u32>()).map(|words| {
        let mut scratch = [0_u8; size_of::<u32>()];
        scratch[..words.len()].copy_from_slice(words);
        u32::from_le_bytes(scratch)
    });

    let mut idx = 0;
    for word in words {
        if idx == len {
            ral::write_reg!(flexspi, flexspi, INTR, IPTXWE: 1);
            idx = 0;
        }
        while ral::read_reg!(flexspi, flexspi, INTR, IPTXWE != 1) {}
        ral::write_reg!(flexspi, flexspi, TFDR[idx], word);
        idx += 1;
    }

    ral::write_reg!(flexspi, flexspi, INTR, IPTXWE: 1);
}

/// Receive bytes from the flash part.
///
/// Returns as soon as data is received into your buffer.
/// When this returns, you must wait for the IP command to complete
/// then clear the RX FIFO.
pub fn receive_bytes(flexspi: flexspi::Instance, data: &mut [u8]) {
    let fifo = FifoReader::new(flexspi, data.len());
    let words = data.chunks_mut(size_of::<u32>());

    for (src, dst) in fifo.zip(words) {
        let src = src.to_le_bytes();
        dst.copy_from_slice(&src[..dst.len()]);
    }

    ral::write_reg!(flexspi, flexspi, INTR, IPRXWA: 1);
}

/// Wait for an IP command to complete.
///
/// This will clear the flag.
pub fn wait_for_ip_cmd_done(flexspi: flexspi::Instance) {
    while ral::read_reg!(flexspi, flexspi, INTR, IPCMDDONE != 1) {}
    ral::write_reg!(flexspi, flexspi, INTR, IPCMDDONE: 1);
}

/// Wait for the FlexSPI interface to become idle.
pub fn wait_for_idle(flexspi: flexspi::Instance) {
    while ral::read_reg!(flexspi, flexspi, STS0, ARBIDLE != 1) {}
}

/// Returns the length of the pointed-at array.
const fn array_ptr_len<T, const N: usize>(_: *const [T; N]) -> usize {
    N
}

/// Returns the number of `u32`s that fit in `TFDR`.
///
/// This may be less than the actual FIFO capacity.
const fn tfdr_len(flexspi: flexspi::Instance) -> usize {
    // Safety: Users swear that the FlexSPI instance's
    // pointer is valid for the whole register block.
    array_ptr_len(unsafe { &raw const (*flexspi.as_ptr()).TFDR })
}

/// Iterator over the RX FIFO.
///
/// After you discard the reader, you should wait for the
/// IP command to complete. Then, you should clear the RX
/// FIFO.
///
/// This implementation performs no eager bounds checking.
/// You must make sure that the watermark is valid in the
/// control register.
struct FifoReader {
    flexspi: flexspi::Instance,
    /// The number of bytes buffered in the FIFO per
    /// the watermark level.
    watermark: usize,
    /// The total number of bytes to receive.
    bytes: usize,
    /// Our index into the FIFO. Each index represents
    /// four bytes pulled from the FIFO.
    idx: usize,
}

impl FifoReader {
    /// Create an interator that reads `bytes` from the RX FIFO.
    fn new(flexspi: flexspi::Instance, bytes: usize) -> Self {
        let watermark =
            (ral::read_reg!(flexspi, flexspi, IPRXFCR, RXWMRK).saturating_add(1) * 8) as usize;
        Self {
            flexspi,
            watermark,
            bytes,
            idx: 0,
        }
    }

    /// Spin until there is data available in the FIFO.
    ///
    /// We're supposed to pivot our polling when we have less than
    /// a watermark of data remaining. Some early testing showed this
    /// isn't strictly necessary, at least on some IP blocks.
    /// Nevertheless, we follow the guidance here. (The DMA path can't
    /// do this, but we have to...?)
    fn wait_for_watermark(&self) {
        if self.bytes >= self.watermark {
            while ral::read_reg!(flexspi, self.flexspi, INTR, IPRXWA != 1) {}
        } else {
            while ral::read_reg!(flexspi, self.flexspi, IPRXFSTS, FILL == 0) {}
        }
    }
}

impl Iterator for FifoReader {
    type Item = u32;

    fn next(&mut self) -> Option<Self::Item> {
        if self.bytes == 0 {
            return None;
        }

        self.wait_for_watermark();

        let word = ral::read_reg!(flexspi, self.flexspi, RFDR[self.idx]);
        self.idx += 1;
        self.bytes = self.bytes.saturating_sub(size_of::<u32>());

        if self.idx * 4 >= self.watermark {
            self.idx = 0;
            ral::write_reg!(flexspi, self.flexspi, INTR, IPRXWA: 1);
        }

        Some(word)
    }
}

pub trait ImxrtFlashAlgorithm: 'static {
    const FLASH_CAPACITY_BYTES: usize;
    const FLASH_SECTOR_SIZE_BYTES: usize;
    const FLASH_PAGE_SIZE_BYTES: usize;

    fn initialize(flexspi: flexspi::Instance);
    fn deinitialize(flexspi: flexspi::Instance) {
        crate::flash::reset(flexspi);
    }
}