use core::{marker::PhantomData, ops::Index};

use alloc::sync::Arc;
use device_api::{
    clock::{ClockController, ClockHandle, Hertz, ResetController, ResetHandle},
    device::{Device, DeviceInitContext},
};
use device_tree::{
    driver::{
        device_tree_driver, lookup_phandle, DeviceTreeClockController, DeviceTreeResetController,
        DeviceTreeSyscon, Node, ProbeContext,
    },
    DeviceTreePropertyRead, TProp,
};
use libk_mm::{address::PhysicalAddress, device::DeviceMemoryIoMut};
use libk_util::sync::IrqSafeSpinlock;
use yggdrasil_abi::error::Error;

// SYSCRG clocks

const SYSCRG_CPU_ROOT: usize = 0x000 / 4;
const SYSCRG_CPU_CORE: usize = 0x004 / 4;
const SYSCRG_CPU_BUS: usize = 0x008 / 4;
const SYSCRG_GPU_ROOT: usize = 0x00C / 4;
const SYSCRG_PERH_ROOT: usize = 0x010 / 4;
const SYSCRG_BUS_ROOT: usize = 0x014 / 4;
const SYSCRG_NOCSTG_BUS: usize = 0x018 / 4;
const SYSCRG_AXI_CFG0: usize = 0x01C / 4;
const SYSCRG_STG_AXIAHB: usize = 0x020 / 4;
const SYSCRG_AHB0: usize = 0x024 / 4;
const SYSCRG_AHB1: usize = 0x028 / 4;
const SYSCRG_APB_BUS: usize = 0x02C / 4;
const SYSCRG_APB0: usize = 0x030 / 4;
// usb
const SYSCRG_USB_125M: usize = 0x17C / 4;
const SYSCRG_NOC_BUS_STG_AXI: usize = 0x180 / 4;
// misc
const SYSCRG_IOMUX_APB: usize = 0x1C0 / 4;
// gmac1
const SYSCRG_GMAC1_AHB: usize = 0x184 / 4;
const SYSCRG_GMAC1_AXI: usize = 0x188 / 4;
const SYSCRG_GMAC_SRC: usize = 0x18C / 4;
const SYSCRG_GMAC1_GTXCLK: usize = 0x190 / 4;
const SYSCRG_GMAC1_RMII_RTX: usize = 0x194 / 4;
const SYSCRG_GMAC1_PTP: usize = 0x198 / 4;
const SYSCRG_GMAC1_RX: usize = 0x19C / 4;
const SYSCRG_GMAC1_RX_INV: usize = 0x1A0 / 4;
const SYSCRG_GMAC1_TX: usize = 0x1A4 / 4;
const SYSCRG_GMAC1_TX_INV: usize = 0x1A8 / 4;
const SYSCRG_GMAC1_GTXC: usize = 0x1AC / 4;
// gmac0
const SYSCRG_GMAC0_GTXCLK: usize = 0x1B0 / 4;
const SYSCRG_GMAC0_PTP: usize = 0x1B4 / 4;
const SYSCRG_GMAC_PHY: usize = 0x1B8 / 4;
const SYSCRG_GMAC0_GTXC: usize = 0x1BC / 4;
// TODO ...
// uart
const SYSCRG_UART0_APB: usize = 0x244 / 4;
const SYSCRG_UART0_CORE: usize = 0x248 / 4;
// jtag
const SYSCRG_JTAG_CERTIFICATION_TRNG: usize = 0x2F4 / 4;
const SYSCRG_CLOCK_COUNT: usize = SYSCRG_JTAG_CERTIFICATION_TRNG + 1;

const SYSCRG_OSC: usize = SYSCRG_CLOCK_COUNT;
const SYSCRG_GMAC1_RMII_REFIN: usize = SYSCRG_CLOCK_COUNT + 1;
const SYSCRG_GMAC1_RGMII_RXIN: usize = SYSCRG_CLOCK_COUNT + 2;
#[allow(unused)]
const SYSCRG_I2STX_BCLK_EXT: usize = SYSCRG_CLOCK_COUNT + 3;
#[allow(unused)]
const SYSCRG_I2STX_LRCK_EXT: usize = SYSCRG_CLOCK_COUNT + 4;
#[allow(unused)]
const SYSCRG_I2SRX_BCLK_EXT: usize = SYSCRG_CLOCK_COUNT + 5;
#[allow(unused)]
const SYSCRG_I2SRX_LRCK_EXT: usize = SYSCRG_CLOCK_COUNT + 6;
#[allow(unused)]
const SYSCRG_TDM_EXT: usize = SYSCRG_CLOCK_COUNT + 7;
#[allow(unused)]
const SYSCRG_MCLK_EXT: usize = SYSCRG_CLOCK_COUNT + 8;
const SYSCRG_PLL0_OUT: usize = SYSCRG_CLOCK_COUNT + 9;
const SYSCRG_PLL1_OUT: usize = SYSCRG_CLOCK_COUNT + 10;
const SYSCRG_PLL2_OUT: usize = SYSCRG_CLOCK_COUNT + 11;

// AONCRG clocks

const AONCRG_OSC_DIV4: usize = 0x00 / 4;
const AONCRG_APB_FUNC: usize = 0x04 / 4;
// gmac0
const AONCRG_GMAC0_AHB: usize = 0x08 / 4;
const AONCRG_GMAC0_AXI: usize = 0x0C / 4;
const AONCRG_GMAC0_RMII_RTX: usize = 0x10 / 4;
const AONCRG_GMAC0_TX: usize = 0x14 / 4;
const AONCRG_GMAC0_TX_INV: usize = 0x18 / 4;
const AONCRG_GMAC0_RX: usize = 0x1C / 4;
const AONCRG_GMAC0_RX_INV: usize = 0x20 / 4;
// otpc
const AONCRG_OTPC_APB: usize = 0x24 / 4;
// rtc
const AONCRG_RTC_APB: usize = 0x28 / 4;
const AONCRG_RTC_INTERNAL: usize = 0x2C / 4;
const AONCRG_RTC_32K: usize = 0x30 / 4;
const AONCRG_RTC_CAL: usize = 0x34 / 4;

const AONCRG_CLOCK_COUNT: usize = AONCRG_RTC_CAL + 1;

const AONCRG_OSC: usize = AONCRG_CLOCK_COUNT;
const AONCRG_GMAC0_RMII_REFIN: usize = AONCRG_CLOCK_COUNT + 1;
const AONCRG_GMAC0_RGMII_RXIN: usize = AONCRG_CLOCK_COUNT + 2;
const AONCRG_STG_AXIAHB: usize = AONCRG_CLOCK_COUNT + 3;
const AONCRG_APB_BUS: usize = AONCRG_CLOCK_COUNT + 4;
const AONCRG_GMAC0_GTXCLK: usize = AONCRG_CLOCK_COUNT + 5;
const AONCRG_RTC_OSC: usize = AONCRG_CLOCK_COUNT + 6;

// STGCRG clocks

// hifi4
const STGCRG_HIFI4_CLK_CORE: usize = 0x00 / 4;
// usb
const STGCRG_USB_APB: usize = 0x04 / 4;
const STGCRG_USB_UTMI_APB: usize = 0x08 / 4;
const STGCRG_USB_AXI: usize = 0x0C / 4;
const STGCRG_USB_LPM: usize = 0x10 / 4;
const STGCRG_USB_STB: usize = 0x14 / 4;
const STGCRG_USB_APP_125: usize = 0x18 / 4;
const STGCRG_USB_REF: usize = 0x1C / 4;
const STGCRG_PCIE0_AXI_MST0: usize = 0x20 / 4;
const STGCRG_PCIE0_APB: usize = 0x24 / 4;
const STGCRG_PCIE0_TL: usize = 0x28 / 4;
const STGCRG_PCIE1_AXI_MST0: usize = 0x2C / 4;
const STGCRG_PCIE1_APB: usize = 0x30 / 4;
const STGCRG_PCIE1_TL: usize = 0x34 / 4;
const STGCRG_PCIE1_SLV_DEC: usize = 0x38 / 4;
const STGCRG_SECURITY_HCLK: usize = 0x3C / 4;
const STGCRG_SECURITY_MISC_AHB: usize = 0x40 / 4;
const STGCRG_MTRX_GRP0: usize = 0x44 / 4;
const STGCRG_MTRX_GRP0_BUS: usize = 0x48 / 4;
const STGCRG_MTRX_GRP0_STG: usize = 0x4C / 4;
const STGCRG_MTRX_GRP1: usize = 0x50 / 4;
const STGCRG_MTRX_GRP1_BUS: usize = 0x54 / 4;
const STGCRG_MTRX_GRP1_STG: usize = 0x58 / 4;
const STGCRG_MTRX_GRP1_HIFI: usize = 0x5C / 4;
const STGCRG_E24_RTC: usize = 0x60 / 4;
const STGCRG_E24_CORE: usize = 0x64 / 4;
const STGCRG_E24_DBG: usize = 0x68 / 4;
const STGCRG_DMA1P_AXI: usize = 0x6C / 4;
const STGCRG_DMA1P_AHB: usize = 0x70 / 4;

const STGCRG_CLOCK_COUNT: usize = STGCRG_DMA1P_AHB + 1;

const STGCRG_OSC: usize = STGCRG_CLOCK_COUNT;
const STGCRG_HIFI4_CORE: usize = STGCRG_CLOCK_COUNT + 1;
const STGCRG_STG_AXIAHB: usize = STGCRG_CLOCK_COUNT + 2;
const STGCRG_USB_125M: usize = STGCRG_CLOCK_COUNT + 3;
const STGCRG_CPU_BUS: usize = STGCRG_CLOCK_COUNT + 4;
const STGCRG_HIFI4_AXI: usize = STGCRG_CLOCK_COUNT + 5;
const STGCRG_NOCSTG_BUS: usize = STGCRG_CLOCK_COUNT + 6;
const STGCRG_APB_BUS: usize = STGCRG_CLOCK_COUNT + 7;

#[allow(unused)]
enum ClockDef {
    Mux(&'static [usize]),
    Div(u32, usize),
    MuxDiv(u32, &'static [usize]),
    Gate(usize),
    GateDiv(u32, usize),
    GateMux(&'static [usize]),
    Inv(usize),
}

trait ClockDefinition {
    const NAME: &'static str;
    const CLOCKS: &'static [Option<(&'static str, ClockDef)>];
}

const SYSCRG_CLOCKS: &'static [Option<(&'static str, ClockDef)>] = &const {
    use ClockDef::*;

    let mut t = [const { None }; SYSCRG_CLOCK_COUNT];

    t[SYSCRG_CPU_ROOT] = Some(("clk_cpu_root", Mux(&[SYSCRG_OSC, SYSCRG_PLL0_OUT])));
    t[SYSCRG_CPU_CORE] = Some(("clk_cpu_core", Div(7, SYSCRG_CPU_ROOT)));
    t[SYSCRG_CPU_BUS] = Some(("clk_cpu_bus", Div(2, SYSCRG_CPU_CORE)));
    t[SYSCRG_GPU_ROOT] = Some(("clk_gpu_root", Mux(&[SYSCRG_PLL2_OUT, SYSCRG_PLL1_OUT])));
    t[SYSCRG_PERH_ROOT] = Some((
        "clk_perh_root",
        MuxDiv(2, &[SYSCRG_PLL0_OUT, SYSCRG_PLL2_OUT]),
    ));
    t[SYSCRG_BUS_ROOT] = Some(("clk_bus_root", Mux(&[SYSCRG_OSC, SYSCRG_PLL2_OUT])));
    t[SYSCRG_NOCSTG_BUS] = Some(("clk_nocstg_bus", Div(3, SYSCRG_BUS_ROOT)));
    t[SYSCRG_AXI_CFG0] = Some(("clk_axi_cfg0", Div(3, SYSCRG_BUS_ROOT)));
    t[SYSCRG_STG_AXIAHB] = Some(("clk_stg_axiahb", Div(2, SYSCRG_AXI_CFG0)));
    t[SYSCRG_AHB0] = Some(("clk_ahb0", Gate(SYSCRG_STG_AXIAHB)));
    t[SYSCRG_AHB1] = Some(("clk_ahb1", Gate(SYSCRG_STG_AXIAHB)));
    t[SYSCRG_APB_BUS] = Some(("clk_apb_bus", Div(8, SYSCRG_STG_AXIAHB)));
    t[SYSCRG_APB0] = Some(("clk_apb0", Gate(SYSCRG_APB_BUS)));
    // usb
    t[SYSCRG_USB_125M] = Some(("clk_usb_125m", Div(15, SYSCRG_PLL0_OUT)));
    t[SYSCRG_NOC_BUS_STG_AXI] = Some(("clk_noc_bus_stg_axi", Gate(SYSCRG_NOCSTG_BUS)));
    // misc
    t[SYSCRG_IOMUX_APB] = Some(("clk_iomux_apb", Gate(SYSCRG_APB_BUS)));
    // gmac1
    t[SYSCRG_GMAC1_AHB] = Some(("clk_gmac1_ahb", Gate(SYSCRG_AHB0)));
    t[SYSCRG_GMAC1_AXI] = Some(("clk_gmac1_axi", Gate(SYSCRG_STG_AXIAHB)));
    t[SYSCRG_GMAC_SRC] = Some(("clk_gmac_src", Div(7, SYSCRG_PLL0_OUT)));
    t[SYSCRG_GMAC1_GTXCLK] = Some(("clk_gmac1_gtxclk", Div(15, SYSCRG_PLL0_OUT)));
    t[SYSCRG_GMAC1_RMII_RTX] = Some(("clk_gmac1_rmii_rtx", Div(30, SYSCRG_GMAC1_RMII_REFIN)));
    t[SYSCRG_GMAC1_PTP] = Some(("clk_gmac1_ptp", GateDiv(31, SYSCRG_GMAC_SRC)));
    t[SYSCRG_GMAC1_RX] = Some((
        "clk_gmac1_rx",
        Mux(&[SYSCRG_GMAC1_RGMII_RXIN, SYSCRG_GMAC1_RMII_RTX]),
    ));
    t[SYSCRG_GMAC1_RX_INV] = Some(("clk_gmac1_rx_inv", Inv(SYSCRG_GMAC1_RX)));
    t[SYSCRG_GMAC1_TX] = Some((
        "clk_gmac1_tx",
        GateMux(&[SYSCRG_GMAC1_GTXCLK, SYSCRG_GMAC1_RMII_RTX]),
    ));
    t[SYSCRG_GMAC1_TX_INV] = Some(("clk_gmac1_tx_inv", Inv(SYSCRG_GMAC1_TX)));
    t[SYSCRG_GMAC1_GTXC] = Some(("clk_gmac1_gtxc", Gate(SYSCRG_GMAC1_GTXCLK)));
    // gmac0
    t[SYSCRG_GMAC0_GTXCLK] = Some(("clk_gmac0_gtxclk", GateDiv(15, SYSCRG_PLL0_OUT)));
    t[SYSCRG_GMAC0_PTP] = Some(("clk_gmac0_ptp", GateDiv(31, SYSCRG_GMAC_SRC)));
    t[SYSCRG_GMAC_PHY] = Some(("clk_gmac_phy", GateDiv(31, SYSCRG_GMAC_SRC)));
    t[SYSCRG_GMAC0_GTXC] = Some(("clk_gmac0_gtxc", Gate(SYSCRG_GMAC0_GTXCLK)));

    // TODO ...

    // uart
    t[SYSCRG_UART0_APB] = Some(("clk_uart0_apb", Gate(SYSCRG_APB0)));
    t[SYSCRG_UART0_CORE] = Some(("clk_uart0_core", Gate(SYSCRG_OSC)));

    t
};

const AONCRG_CLOCKS: &'static [Option<(&'static str, ClockDef)>] = &const {
    use ClockDef::*;

    let mut t = [const { None }; AONCRG_CLOCK_COUNT];

    t[AONCRG_OSC_DIV4] = Some(("clk_osc_div4", Div(4, AONCRG_OSC)));
    t[AONCRG_APB_FUNC] = Some(("clk_apb_func", Mux(&[AONCRG_OSC_DIV4, AONCRG_OSC])));
    // gmac0
    t[AONCRG_GMAC0_AHB] = Some(("clk_gmac0_ahb", Gate(AONCRG_STG_AXIAHB)));
    t[AONCRG_GMAC0_AXI] = Some(("clk_gmac0_axi", Gate(AONCRG_STG_AXIAHB)));
    t[AONCRG_GMAC0_RMII_RTX] = Some(("clk_gmac0_rmii_rtx", Div(30, AONCRG_GMAC0_RMII_REFIN)));
    t[AONCRG_GMAC0_TX] = Some((
        "clk_gmac0_tx",
        GateMux(&[AONCRG_GMAC0_GTXCLK, AONCRG_GMAC0_RMII_RTX]),
    ));
    t[AONCRG_GMAC0_TX_INV] = Some(("clk_gmac0_tx_inv", Inv(AONCRG_GMAC0_TX)));
    t[AONCRG_GMAC0_RX] = Some((
        "clk_gmac0_rx",
        Mux(&[AONCRG_GMAC0_RGMII_RXIN, AONCRG_GMAC0_RMII_RTX]),
    ));
    t[AONCRG_GMAC0_RX_INV] = Some(("clk_gmac0_rx_inv", Inv(AONCRG_GMAC0_RX)));
    // otpc
    t[AONCRG_OTPC_APB] = Some(("clk_otpc_apb", Gate(AONCRG_APB_BUS)));
    // rtc
    t[AONCRG_RTC_APB] = Some(("clk_rtc_apb", Gate(AONCRG_APB_BUS)));
    t[AONCRG_RTC_INTERNAL] = Some(("clk_rtc_internal", Div(1022, AONCRG_OSC)));
    t[AONCRG_RTC_32K] = Some(("clk_rtc_32k", Mux(&[AONCRG_RTC_OSC, AONCRG_RTC_INTERNAL])));
    t[AONCRG_RTC_CAL] = Some(("clk_rtc_cal", Gate(AONCRG_OSC)));

    t
};

const STGCRG_CLOCKS: &'static [Option<(&'static str, ClockDef)>] = &const {
    use ClockDef::*;

    let mut t = [const { None }; STGCRG_CLOCK_COUNT];

    // hifi4
    t[STGCRG_HIFI4_CLK_CORE] = Some(("clk_hifi4_clk_core", Gate(STGCRG_HIFI4_CORE)));
    // usb
    t[STGCRG_USB_APB] = Some(("clk_usb_apb", Gate(STGCRG_APB_BUS)));
    t[STGCRG_USB_UTMI_APB] = Some(("clk_usb_utmi_apb", Gate(STGCRG_APB_BUS)));
    t[STGCRG_USB_AXI] = Some(("clk_usb_axi", Gate(STGCRG_STG_AXIAHB)));
    t[STGCRG_USB_LPM] = Some(("clk_usb_lpm", GateDiv(2, STGCRG_OSC)));
    t[STGCRG_USB_STB] = Some(("clk_usb_stb", GateDiv(4, STGCRG_OSC)));
    t[STGCRG_USB_APP_125] = Some(("clk_usb_app_125", Gate(STGCRG_USB_125M)));
    t[STGCRG_USB_REF] = Some(("clk_usb_ref", Div(2, STGCRG_OSC)));
    // pcie
    t[STGCRG_PCIE0_AXI_MST0] = Some(("clk_pcie0_axi_mst0", Gate(STGCRG_STG_AXIAHB)));
    t[STGCRG_PCIE0_APB] = Some(("clk_pcie0_apb", Gate(STGCRG_APB_BUS)));
    t[STGCRG_PCIE0_TL] = Some(("clk_pcie0_tl", Gate(STGCRG_STG_AXIAHB)));
    t[STGCRG_PCIE1_AXI_MST0] = Some(("clk_pcie1_axi_mst0", Gate(STGCRG_STG_AXIAHB)));
    t[STGCRG_PCIE1_APB] = Some(("clk_pcie1_apb", Gate(STGCRG_APB_BUS)));
    t[STGCRG_PCIE1_TL] = Some(("clk_pcie1_tl", Gate(STGCRG_STG_AXIAHB)));
    t[STGCRG_PCIE1_SLV_DEC] = Some(("clk_pcie1_slv_dec", Gate(STGCRG_STG_AXIAHB)));
    // security
    t[STGCRG_SECURITY_HCLK] = Some(("clk_security_hclk", Gate(STGCRG_STG_AXIAHB)));
    t[STGCRG_SECURITY_MISC_AHB] = Some(("clk_security_misc_ahb", Gate(STGCRG_STG_AXIAHB)));
    // mtrx
    t[STGCRG_MTRX_GRP0] = Some(("clk_mtrx_grp0", Gate(STGCRG_CPU_BUS)));
    t[STGCRG_MTRX_GRP0_BUS] = Some(("clk_mrtx_grp0_bus", Gate(STGCRG_NOCSTG_BUS)));
    t[STGCRG_MTRX_GRP0_STG] = Some(("clk_mtrx_grp0_stg", Gate(STGCRG_STG_AXIAHB)));
    t[STGCRG_MTRX_GRP1] = Some(("clk_mtrx_grp1", Gate(STGCRG_CPU_BUS)));
    t[STGCRG_MTRX_GRP1_BUS] = Some(("clk_mtrx_grp1_bus", Gate(STGCRG_NOCSTG_BUS)));
    t[STGCRG_MTRX_GRP1_STG] = Some(("clk_mtrx_grp1_stg", Gate(STGCRG_STG_AXIAHB)));
    t[STGCRG_MTRX_GRP1_HIFI] = Some(("clk_mtrx_grp1_hifi", Gate(STGCRG_HIFI4_AXI)));
    // e24
    t[STGCRG_E24_RTC] = Some(("clk_e24_rtc", GateDiv(24, STGCRG_OSC)));
    t[STGCRG_E24_CORE] = Some(("clk_e24_core", Gate(STGCRG_STG_AXIAHB)));
    t[STGCRG_E24_DBG] = Some(("clk_e24_dbg", Gate(STGCRG_STG_AXIAHB)));
    // dma1p
    t[STGCRG_DMA1P_AXI] = Some(("clk_dma1p_axi", Gate(STGCRG_STG_AXIAHB)));
    t[STGCRG_DMA1P_AHB] = Some(("clk_dma1p_ahb", Gate(STGCRG_STG_AXIAHB)));

    t
};

struct Crg<C, P> {
    clock_mapping: IrqSafeSpinlock<DeviceMemoryIoMut<'static, [u32]>>,
    reset_mapping: IrqSafeSpinlock<DeviceMemoryIoMut<'static, [u32]>>,
    parents: P,
    _pd: PhantomData<C>,
}

struct Pll {
    syscon: Arc<dyn DeviceTreeSyscon>,
    clk_osc: ClockHandle,
}

struct PllDef {
    fbdiv_offset: usize,
    pd_offset: usize,
    frac_offset: usize,
    prediv_offset: usize,

    dacpd_bit: u32,
    dsmpd_bit: u32,

    fbdiv_shift: usize,
}

#[derive(Debug)]
struct PllRegs {
    dacpd: bool,
    dsmpd: bool,
    fbdiv: u32,
    #[allow(unused)]
    frac: u32,
    postdiv1: u32,
    prediv: u32,
}

unsafe impl<C, P> Send for Crg<C, P> {}
unsafe impl<C, P> Sync for Crg<C, P> {}

impl<C: ClockDefinition, P: Index<usize, Output = ClockHandle>> Crg<C, P> {
    unsafe fn map(base: PhysicalAddress, parents: P) -> Result<Arc<Self>, Error> {
        let clock_mapping =
            DeviceMemoryIoMut::map_slice(base, C::CLOCKS.len(), Default::default())?;
        let reset_mapping =
            DeviceMemoryIoMut::map_slice(base.add(C::CLOCKS.len() * 4), 16, Default::default())?;
        Ok(Arc::new(Self {
            clock_mapping: IrqSafeSpinlock::new(clock_mapping),
            reset_mapping: IrqSafeSpinlock::new(reset_mapping),
            parents,
            _pd: PhantomData,
        }))
    }

    fn clock(&self, index: usize) -> Result<&(&str, ClockDef), Error> {
        C::CLOCKS
            .get(index)
            .ok_or(Error::DoesNotExist)?
            .as_ref()
            .ok_or(Error::DoesNotExist)
            .inspect_err(|_| {
                log::error!(
                    "{}: clock {:#x} does not exist/is not implemented",
                    C::NAME,
                    index
                )
            })
    }

    fn read_clock_reg(&self, index: usize) -> u32 {
        self.clock_mapping.lock()[index]
    }

    fn enable_clock_inner(&self, index: usize) -> Result<(), Error> {
        if index >= C::CLOCKS.len() {
            return self.parents[index - C::CLOCKS.len()].enable();
        }

        let (clk_name, clk_def) = self.clock(index)?;

        let (parent, write_gate) = match clk_def {
            ClockDef::GateDiv(_, parent) | ClockDef::Gate(parent) => (*parent, true),
            ClockDef::GateMux(parents) => (parents[0], true),

            ClockDef::Inv(parent) | ClockDef::Div(_, parent) => (*parent, false),
            ClockDef::Mux(parents) | ClockDef::MuxDiv(_, parents) => (parents[0], false),
        };

        // Enable parent
        self.enable_clock_inner(parent)?;

        // Enable clock
        if write_gate {
            let mut lock = self.clock_mapping.lock();
            let ptr = &raw mut lock[index];
            unsafe {
                let val = ptr.read_volatile();
                if val & (1 << 31) == 0 {
                    log::info!("{}: enable clock {:?}", C::NAME, clk_name);
                    ptr.write_volatile(val | (1 << 31));
                }
            }
        }

        Ok(())
    }

    fn clock_rate_inner(&self, index: usize) -> Result<Hertz, Error> {
        if index >= C::CLOCKS.len() {
            return self.parents[index - C::CLOCKS.len()].rate();
        }

        let (_, clk_def) = self.clock(index)?;

        let reg = self.read_clock_reg(index);
        let reg_div = reg & 0xFFFFFF;

        let (parent, div) = match clk_def {
            &ClockDef::Gate(parent) | &ClockDef::Inv(parent) => (parent, 1),
            &ClockDef::Div(_max, parent) => (parent, reg_div),
            // TODO read actual parent for muxes
            &ClockDef::GateDiv(_max, parent) => (parent, reg_div),
            ClockDef::GateMux(parents) | ClockDef::Mux(parents) => (parents[0], 1),
            ClockDef::MuxDiv(_max, parents) => (parents[0], reg_div),
        };

        let parent_rate = self.clock_rate_inner(parent)?;

        Ok(parent_rate / div)
    }

    fn set_reset_asserted_inner(&self, index: usize, asserted: bool) -> Result<(), Error> {
        let reg_index = index / 32;
        let reg_shift = index % 32;

        let mut lock = self.reset_mapping.lock();
        let ptr = &raw mut lock[reg_index];
        unsafe {
            let val = ptr.read_volatile();
            if asserted {
                ptr.write_volatile(val | (1 << reg_shift));
            } else {
                ptr.write_volatile(val & !(1 << reg_shift));
            }
        }

        Ok(())
    }
}

impl<C: ClockDefinition, P> Device for Crg<C, P> {
    fn display_name(&self) -> &str {
        C::NAME
    }
}

impl<C: ClockDefinition, P: Index<usize, Output = ClockHandle>> ClockController for Crg<C, P> {
    fn clock_rate(&self, clock: Option<u32>) -> Result<Hertz, Error> {
        let index = clock.ok_or(Error::InvalidArgument)? as usize;
        self.clock_rate_inner(index)
    }

    fn set_clock_rate(&self, _clock: Option<u32>, _rate: Hertz) -> Result<Hertz, Error> {
        todo!("Clock rate update not supported for jh71x0 CRGs")
    }

    fn enable_clock(&self, clock: Option<u32>) -> Result<(), Error> {
        let index = clock.ok_or(Error::InvalidArgument)? as usize;
        self.enable_clock_inner(index)
    }

    fn disable_clock(&self, _clock: Option<u32>) -> Result<(), Error> {
        todo!("Clock disable not supported for jh71x0 CRGs")
    }
}

impl<C: ClockDefinition, P: Index<usize, Output = ClockHandle>> ResetController for Crg<C, P> {
    fn assert_reset(&self, reset: Option<u32>) -> Result<(), Error> {
        let reset = reset.ok_or(Error::InvalidArgument)? as usize;
        self.set_reset_asserted_inner(reset, true)
    }

    fn deassert_reset(&self, reset: Option<u32>) -> Result<(), Error> {
        let reset = reset.ok_or(Error::InvalidArgument)? as usize;
        self.set_reset_asserted_inner(reset, false)
    }
}

impl<C: ClockDefinition + 'static, P: Index<usize, Output = ClockHandle> + 'static>
    DeviceTreeClockController for Crg<C, P>
{
    fn map_clock(self: Arc<Self>, property: &TProp, offset: usize) -> Option<(ClockHandle, usize)> {
        let clock = property.read_cell(offset, 1)? as u32;
        Some((
            ClockHandle {
                clock: Some(clock),
                parent: self.clone(),
            },
            1,
        ))
    }
}

impl<C: ClockDefinition + 'static, P: Index<usize, Output = ClockHandle> + 'static>
    DeviceTreeResetController for Crg<C, P>
{
    fn map_reset(self: Arc<Self>, property: &TProp, offset: usize) -> Option<(ResetHandle, usize)> {
        let reset = property.read_cell(offset, 1)? as u32;
        Some((
            ResetHandle {
                reset: Some(reset),
                parent: self.clone(),
            },
            1,
        ))
    }
}

impl Pll {
    const PLLS: &[PllDef] = &[
        PllDef {
            pd_offset: 0x18,
            fbdiv_offset: 0x1C,
            frac_offset: 0x20,
            prediv_offset: 0x24,

            dacpd_bit: 1 << 24,
            dsmpd_bit: 1 << 25,

            fbdiv_shift: 0,
        },
        PllDef {
            pd_offset: 0x24,
            fbdiv_offset: 0x24,
            frac_offset: 0x28,
            prediv_offset: 0x2C,

            dacpd_bit: 1 << 15,
            dsmpd_bit: 1 << 16,

            fbdiv_shift: 17,
        },
        PllDef {
            pd_offset: 0x2C,
            fbdiv_offset: 0x2C,
            frac_offset: 0x30,
            prediv_offset: 0x34,

            dacpd_bit: 1 << 15,
            dsmpd_bit: 1 << 16,

            fbdiv_shift: 17,
        },
    ];

    const FBDIV_MASK: u32 = 0xFFF;
    const FRAC_MASK: u32 = 0xFFFFFF;
    const POSTDIV1_MASK: u32 = 0x3;
    const PREDIV_MASK: u32 = 0x3F;

    const POSTDIV1_SHIFT: usize = 28;
    const PREDIV_SHIFT: usize = 0;
    const FRAC_SHIFT: usize = 0;

    fn read_pll_regs(&self, index: usize) -> Result<PllRegs, Error> {
        let pll = &Self::PLLS[index];

        let val = self.syscon.read_register(pll.pd_offset)?;
        let dacpd = val & pll.dacpd_bit != 0;
        let dsmpd = val & pll.dsmpd_bit != 0;

        let val = self.syscon.read_register(pll.fbdiv_offset)?;
        let fbdiv = (val >> pll.fbdiv_shift) & Self::FBDIV_MASK;

        let val = self.syscon.read_register(pll.frac_offset)?;
        let frac = (val >> Self::FRAC_SHIFT) & Self::FRAC_MASK;
        let postdiv1 = (val >> Self::POSTDIV1_SHIFT) & Self::POSTDIV1_MASK;

        let val = self.syscon.read_register(pll.prediv_offset)?;
        let prediv = (val >> Self::PREDIV_SHIFT) & Self::PREDIV_MASK;

        Ok(PllRegs {
            dacpd,
            dsmpd,

            fbdiv,
            frac,
            postdiv1,
            prediv,
        })
    }
}

impl Device for Pll {
    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        // Dump PLL rates
        for i in 0..3 {
            let regs = self.read_pll_regs(i as usize)?;
            let rate = self.clock_rate(Some(i))?;
            log::info!("PLL{i} rate: {rate}");
            log::info!("  {regs:?}");
        }

        Ok(())
    }

    fn display_name(&self) -> &str {
        "jh7110-pll"
    }
}

impl ClockController for Pll {
    fn set_clock_rate(&self, _clock: Option<u32>, _rate: Hertz) -> Result<Hertz, Error> {
        todo!("PLL rate configuration not yet implemented")
    }

    fn clock_rate(&self, clock: Option<u32>) -> Result<Hertz, Error> {
        let index = clock.ok_or(Error::InvalidArgument)? as usize;
        if index >= 3 {
            return Err(Error::InvalidArgument);
        }

        let osc = self.clk_osc.rate()?;
        let pll = self.read_pll_regs(index)?;

        if pll.dacpd && pll.dsmpd {
            // Integer mode
            Ok((osc * pll.fbdiv) / (pll.prediv << pll.postdiv1))
        } else if !pll.dacpd && !pll.dsmpd {
            // Fraction mode
            todo!()
        } else {
            todo!("Invalid PLL dacpd/dsmpd combination")
        }
    }

    fn enable_clock(&self, _clock: Option<u32>) -> Result<(), Error> {
        // Assumed always enabled
        Ok(())
    }

    fn disable_clock(&self, _clock: Option<u32>) -> Result<(), Error> {
        Err(Error::InvalidOperation)
    }
}

impl DeviceTreeClockController for Pll {
    fn map_clock(self: Arc<Self>, property: &TProp, offset: usize) -> Option<(ClockHandle, usize)> {
        let clock = property.read_cell(offset, 1)? as u32;
        Some((
            ClockHandle {
                clock: Some(clock),
                parent: self.clone(),
            },
            1,
        ))
    }
}

device_tree_driver! {
    compatible: ["starfive,jh7110-syscrg"],
    driver: {
        fn probe(&self, node: &Arc<Node>, context: &mut ProbeContext) -> Option<Arc<dyn Device>> {
            struct Syscrg;

            impl ClockDefinition for Syscrg {
                const NAME: &'static str = "jh7110-syscrg";
                const CLOCKS: &'static [Option<(&'static str, ClockDef)>] = SYSCRG_CLOCKS;
            }

            let base = node.map_base(context, 0)?;
            let parents = [
                node.named_clock("osc")?,
                node.named_clock("gmac1_rmii_refin")?,
                node.named_clock("gmac1_rgmii_rxin")?,
                node.named_clock("i2stx_bclk_ext")?,
                node.named_clock("i2stx_lrck_ext")?,
                node.named_clock("i2srx_bclk_ext")?,
                node.named_clock("i2srx_lrck_ext")?,
                node.named_clock("tdm_ext")?,
                node.named_clock("mclk_ext")?,
                node.named_clock("pll0_out")?,
                node.named_clock("pll1_out")?,
                node.named_clock("pll2_out")?
            ];

            let syscrg = unsafe { Crg::<Syscrg, _>::map(base, parents) }.ok()?;

            node.make_clock_controller(syscrg.clone());
            node.make_reset_controller(syscrg.clone());

            Some(syscrg)
        }
    }
}

device_tree_driver! {
    compatible: ["starfive,jh7110-aoncrg"],
    driver: {
        fn probe(&self, node: &Arc<Node>, context: &mut ProbeContext) -> Option<Arc<dyn Device>> {
            struct Aoncrg;

            impl ClockDefinition for Aoncrg {
                const NAME: &'static str = "jh7110-aoncrg";
                const CLOCKS: &'static [Option<(&'static str, ClockDef)>] = AONCRG_CLOCKS;
            }

            let base = node.map_base(context, 0)?;
            let parents = [
                node.named_clock("osc")?,
                node.named_clock("gmac0_rmii_refin")?,
                node.named_clock("gmac0_rgmii_rxin")?,
                node.named_clock("stg_axiahb")?,
                node.named_clock("apb_bus")?,
                node.named_clock("gmac0_gtxclk")?,
                node.named_clock("rtc_osc")?,
            ];

            let aoncrg = unsafe { Crg::<Aoncrg, _>::map(base, parents) }.ok()?;

            node.make_clock_controller(aoncrg.clone());
            node.make_reset_controller(aoncrg.clone());

            Some(aoncrg)
        }
    }
}

device_tree_driver! {
    compatible: ["starfive,jh7110-stgcrg"],
    driver: {
        fn probe(&self, node: &Arc<Node>, context: &mut ProbeContext) -> Option<Arc<dyn Device>> {
            struct Stgcrg;

            impl ClockDefinition for Stgcrg {
                const NAME: &'static str = "jh7110-stgcrg";
                const CLOCKS: &'static [Option<(&'static str, ClockDef)>] = STGCRG_CLOCKS;
            }

            let base = node.map_base(context, 0)?;
            let parents = [
                node.named_clock("osc").unwrap(),
                node.named_clock("hifi4_core").unwrap(),
                node.named_clock("stg_axiahb").unwrap(),
                node.named_clock("usb_125m").unwrap(),
                node.named_clock("cpu_bus").unwrap(),
                node.named_clock("hifi4_axi").unwrap(),
                node.named_clock("nocstg_bus").unwrap(),
                node.named_clock("apb_bus").unwrap()
            ];

            let stgcrg = unsafe { Crg::<Stgcrg, _>::map(base, parents) }.ok()?;

            node.make_clock_controller(stgcrg.clone());
            node.make_reset_controller(stgcrg.clone());

            Some(stgcrg)
        }
    }
}

device_tree_driver! {
    compatible: ["starfive,jh7110-pll"],
    driver: {
        fn probe(&self, node: &Arc<Node>, _context: &mut ProbeContext) -> Option<Arc<dyn Device>> {
            let clk_osc = node.clock(0)?; // 24MHz
            // Make sure parent syscon is probed
            let syscon_phandle = node.parent()?.prop_usize("phandle")? as u32;
            let syscon = lookup_phandle(syscon_phandle, true)?.as_system_controller()?;

            let pll = Arc::new(Pll {
                syscon,
                clk_osc,
            });

            node.make_clock_controller(pll.clone());

            Some(pll)
        }
    }
}
