use core::{marker::PhantomData, num::NonZero};

use imxrt_drivers_ccm_11xx::ral_11xx as ccm;
use imxrt_drivers_flexspi as flexspi;
use imxrt_drivers_gpc_11xx::cpu_mode_ctrl as gpc_cpu;
use imxrt_drivers_pmu_11xx as pmu;
use imxrt_drivers_rtwdog as rtwdog;

const FLEXSPI1_BASE: u32 = 0x3000_0000;

/// Establish the core, bus, and FlexSPI
/// clock frequencies.
fn configure_clocks(
    ccm: ccm::Instance,
    pll: ccm::pll::Instance,
    pmu: pmu::Instance,
    gpc_cpu: gpc_cpu::Instance,
) {
    // Switch the core to something stable before we
    // start changing upstream sources.
    ccm::set_clock_root(
        ccm,
        ccm::ClockRoot::M7,
        const { ccm::mux(ccm::ClockRoot::M7, ccm::ClockSource::Xtal) },
        NO_DIVIDER,
    );
    ccm::set_clock_root(
        ccm,
        ccm::ClockRoot::Bus,
        const { ccm::mux(ccm::ClockRoot::Bus, ccm::ClockSource::Xtal) },
        NO_DIVIDER,
    );
    ccm::set_clock_root(
        ccm,
        ccm::ClockRoot::BusLpsr,
        const { ccm::mux(ccm::ClockRoot::BusLpsr, ccm::ClockSource::Xtal) },
        NO_DIVIDER,
    );
    ccm::set_clock_root(
        ccm,
        ccm::ClockRoot::Flexspi1,
        const { ccm::mux(ccm::ClockRoot::Flexspi1, ccm::ClockSource::Xtal) },
        NO_DIVIDER,
    );

    // Prepare PLL power, GPC setpoints.
    pmu::enable_pll_reference_voltage(pmu, true);
    pmu::set_phy_ldo_setpoints(pmu, u16::MAX);
    pmu::enable_phy_ldo_setpoints(pmu);
    pmu::enable_pll_reference_setpoints(pmu);

    for clock_source in {
        use ccm::ClockSource::*;
        [
            Pll1, Pll1Clk, Pll1Div2, Pll1Div5, ArmPll, ArmPllClk, //
            Pll2, Pll2Clk, Pll2Pfd0, Pll2Pfd1, Pll2Pfd2, Pll2Pfd3, //
            Pll3, Pll3Clk, Pll3Div2, Pll3Pfd0, Pll3Pfd1, Pll3Pfd2, Pll3Pfd3, //
        ]
    } {
        ccm::set_source_setpoints(ccm, clock_source, u16::MAX, 0);
        ccm::enable_source_setpoints(ccm, clock_source).unwrap();
    }

    ccm::pll::enable_sys_pll1_setpoints(pll);
    ccm::pll::enable_arm_pll_setpoints(pll, ARM_PLL_POST_DIV, ARM_PLL_DIV_SELECT);
    ccm::pll::enable_sys_pll2_setpoints(pll);
    ccm::pll::enable_sys_pll3_setpoints(pll);

    gpc_cpu::request_setpoint_transition(gpc_cpu, 1).unwrap();

    ccm::pll::set_pll3_pfd_fracs(
        pll,
        [
            SYS_PLL3_PFD0_DIV,
            SYS_PLL3_PFD1_DIV,
            SYS_PLL3_PFD2_DIV,
            SYS_PLL3_PFD3_DIV,
        ],
    );
    ccm::pll::update_pll3_pfd_fracs(pll, [true, true, true, true]);

    ccm::set_clock_root(
        ccm,
        ccm::ClockRoot::M7,
        const { ccm::mux(ccm::ClockRoot::M7, ccm::ClockSource::ArmPll) },
        M7_DIVIDER,
    );

    ccm::set_clock_root(
        ccm,
        ccm::ClockRoot::Bus,
        const { ccm::mux(ccm::ClockRoot::Bus, ccm::ClockSource::Pll1Div5) },
        BUS_DIVIDER,
    );

    ccm::set_clock_root(
        ccm,
        ccm::ClockRoot::Flexspi1,
        const { ccm::mux(ccm::ClockRoot::Flexspi1, ccm::ClockSource::Pll3Pfd0) },
        FLEXSPI1_DIVIDER,
    );
}

const NO_DIVIDER: NonZero<u8> = NonZero::new(1).unwrap();

const M7_DIVIDER: NonZero<u8> = NO_DIVIDER;
const BUS_DIVIDER: NonZero<u8> = NO_DIVIDER;
const FLEXSPI1_DIVIDER: NonZero<u8> = NonZero::new(2).unwrap();

const ARM_PLL_DIV_SELECT: ccm::pll::ArmPllDivSelect = ccm::pll::ArmPllDivSelect::new(200).unwrap();
const ARM_PLL_POST_DIV: ccm::pll::ArmPllPostDiv = ccm::pll::ArmPllPostDiv::Div4;
const SYS_PLL3_PFD0_DIV: ccm::pll::PfdFrac = ccm::pll::PfdFrac::new(33).unwrap();
const SYS_PLL3_PFD1_DIV: ccm::pll::PfdFrac = ccm::pll::PfdFrac::new(27).unwrap();
const SYS_PLL3_PFD2_DIV: ccm::pll::PfdFrac = ccm::pll::PfdFrac::new(21).unwrap();
const SYS_PLL3_PFD3_DIV: ccm::pll::PfdFrac = ccm::pll::PfdFrac::new(17).unwrap();

pub trait Imxrt11xx: 'static {
    const FLEXSPI1_INSTANCE: flexspi::Instance;
    const CCM_INSTANCE: ccm::Instance;
    const PMU_INSTANCE: pmu::Instance;
    const CCM_PLL_INSTANCE: ccm::pll::Instance;
    const GPC_CPU_INSTANCE: gpc_cpu::Instance;
    const RTWDOG_INSTANCE: rtwdog::Instance;

    const FLEXSPI_FIFO_CAPACITY_BYTES: usize;
}

pub struct Algorithm<C: Imxrt11xx, F: crate::ImxrtFlashAlgorithm>(PhantomData<(C, F)>);

impl<C: Imxrt11xx, F: crate::ImxrtFlashAlgorithm> Algorithm<C, F> {
    pub const fn flash_size_bytes() -> usize {
        F::FLASH_CAPACITY_BYTES
    }
    pub const fn flash_address() -> usize {
        FLEXSPI1_BASE as _
    }
    pub const fn sector_size_bytes() -> usize {
        F::FLASH_SECTOR_SIZE_BYTES
    }
    pub const fn page_size_bytes() -> usize {
        F::FLASH_PAGE_SIZE_BYTES
    }

    pub fn initialize() -> Self {
        rtwdog::disable(C::RTWDOG_INSTANCE);
        configure_clocks(
            C::CCM_INSTANCE,
            C::CCM_PLL_INSTANCE,
            C::PMU_INSTANCE,
            C::GPC_CPU_INSTANCE,
        );
        crate::reset(
            C::FLEXSPI1_INSTANCE,
            F::FLASH_CAPACITY_BYTES / 1024,
            C::FLEXSPI_FIFO_CAPACITY_BYTES,
        );
        F::initialize(C::FLEXSPI1_INSTANCE);
        Algorithm(PhantomData)
    }

    pub fn flash_read(&mut self, address: usize, data: &mut [u8]) {
        crate::flash::read(C::FLEXSPI1_INSTANCE, address, data);
    }

    pub fn flash_erase_sector(&mut self, address: usize) {
        crate::flash::erase_sector(C::FLEXSPI1_INSTANCE, address);
    }

    pub fn flash_write(&mut self, address: usize, data: &[u8]) {
        crate::flash::write(C::FLEXSPI1_INSTANCE, address, data);
    }
}

impl<C: Imxrt11xx, F: crate::ImxrtFlashAlgorithm> flash_algorithm::FlashAlgorithm
    for Algorithm<C, F>
{
    fn new(
        _: u32,
        _: u32,
        _: flash_algorithm::Function,
    ) -> Result<Self, flash_algorithm::ErrorCode> {
        Ok(Self::initialize())
    }

    fn erase_all(&mut self) -> Result<(), flash_algorithm::ErrorCode> {
        crate::flash::erase_chip(C::FLEXSPI1_INSTANCE);
        Ok(())
    }

    fn erase_sector(&mut self, address: u32) -> Result<(), flash_algorithm::ErrorCode> {
        self.flash_erase_sector(address.saturating_sub(FLEXSPI1_BASE) as usize);
        Ok(())
    }

    fn program_page(
        &mut self,
        address: u32,
        data: &[u8],
    ) -> Result<(), flash_algorithm::ErrorCode> {
        self.flash_write(address.saturating_sub(FLEXSPI1_BASE) as usize, data);
        Ok(())
    }

    fn read_flash(
        &mut self,
        address: u32,
        data: &mut [u8],
    ) -> Result<(), flash_algorithm::ErrorCode> {
        self.flash_read(address.saturating_sub(FLEXSPI1_BASE) as usize, data);
        Ok(())
    }
}

impl<C: Imxrt11xx, F: crate::ImxrtFlashAlgorithm> Drop for Algorithm<C, F> {
    fn drop(&mut self) {
        F::deinitialize(C::FLEXSPI1_INSTANCE);
    }
}