//! Starfive JH7110 Always-on CRG
use abi::error::Error;
use alloc::sync::Arc;
use device_api::{
    clock::{ClockController, ClockHandle, ResetController, ResetHandle},
    device::{Device, DeviceInitContext},
};
use device_tree::{
    driver::{
        device_tree_driver, DeviceTreeClockController, DeviceTreeResetController, Node,
        ProbeContext,
    },
    DeviceTreePropertyRead, TProp,
};
use libk_mm::{address::PhysicalAddress, device::DeviceMemoryIo};
use libk_util::{sync::IrqSafeSpinlock, OneTimeInit};
use tock_registers::{
    interfaces::{Readable, Writeable},
    register_structs,
    registers::ReadWrite,
};

const PARENT_OSC: usize = 0;
// const PARENT_GMAC0_RMII_REFIN: usize = 1;
// const PARENT_GMAC0_RGMII_RXIN: usize = 2;
const PARENT_STG_AXIAHB: usize = 3;
// const PARENT_APB_BUS: usize = 4;
// const PARENT_GMAC0_GTXCLK: usize = 5;
// const PARENT_RTC_OSC: usize = 6;

register_structs! {
    #[allow(non_snake_case)]
    Regs {
        (0x000 => CLOCKS: [ReadWrite<u32>; 14]),
        (0x038 => RESETS: ReadWrite<u32>),
        (0x03C => AONCRG_RESET_STATUS: ReadWrite<u32>),
        (0x040 => @END),
    }
}

#[derive(Debug, Clone, Copy)]
enum ClockRegister {
    Gate(u32, ClockParent),
    #[allow(unused)]
    FixedDiv(u32, ClockParent),
    // TODO
    Fixed,
    Unimp,
}

#[derive(Debug, Clone, Copy)]
enum ClockParent {
    Int(u32),
    Ext(usize),
}

/// JH7110 AONCRG driver
pub struct Aoncrg {
    base: PhysicalAddress,
    clock_parents: [ClockHandle; 7],
    mapping: OneTimeInit<IrqSafeSpinlock<DeviceMemoryIo<'static, Regs>>>,
}

impl Aoncrg {
    fn ensure_init(&self) -> Result<&IrqSafeSpinlock<DeviceMemoryIo<'static, Regs>>, Error> {
        self.mapping.or_try_init_with(|| {
            unsafe { DeviceMemoryIo::map(self.base, Default::default()) }.map(IrqSafeSpinlock::new)
        })
    }

    fn clk_enable_int(&self, index: u32) -> Result<(), Error> {
        let (name, reg) = Self::map_clock_index(index)
            .ok_or(Error::InvalidArgument)
            .inspect_err(|_| log::warn!("jh7110-syscrg: undefined clock {:?}", index))?;
        let regs = self.ensure_init()?;

        match reg {
            ClockRegister::Fixed => Ok(()),
            ClockRegister::Gate(bit, parent) => {
                self.clk_enable_parent(parent)?;
                log::info!("jh7110-aoncrg: enable {name:?} @ {index}");
                let lock = regs.lock();
                lock.CLOCKS[index as usize].set(lock.CLOCKS[index as usize].get() | (1 << bit));
                Ok(())
            }
            ClockRegister::Unimp => {
                log::warn!("jh7110-aoncrg: unimplemented clock {name:?}");
                Err(Error::NotImplemented)
            }
            ClockRegister::FixedDiv(_, parent) => {
                self.clk_enable_parent(parent)?;
                Ok(())
            }
        }
    }

    fn clk_enable_parent(&self, parent: ClockParent) -> Result<(), Error> {
        match parent {
            ClockParent::Int(index) => self.clk_enable_int(index),
            ClockParent::Ext(clock) => self.clock_parents[clock].enable(),
        }
    }

    fn map_clock_index(index: u32) -> Option<(&'static str, ClockRegister)> {
        use ClockParent::*;
        use ClockRegister::*;

        const CLOCK_MAP: &[(&'static str, ClockRegister)] = &[
            /* 0 */ ("clk_osc", FixedDiv(4, Ext(PARENT_OSC))),
            /* 1 */ ("clk_aon_apb_func", Unimp),
            /* 2 */ ("clk_gmac5_ahb", Gate(31, Ext(PARENT_STG_AXIAHB))),
            /* 3 */ ("clk_gmac5_axi", Gate(31, Ext(PARENT_STG_AXIAHB))),
            /* 4 */ ("clk_gmac5_rmii_rtx", Unimp),
            /* 5 */ ("clk_gmac5_axi64_tx", Fixed),
            /* 6 */ ("clk_gmac5_axi64_tx_inv", Gate(30, Int(5))),
        ];

        CLOCK_MAP.get(index as usize).copied()
    }
}

impl Device for Aoncrg {
    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        log::warn!("TODO: init jh7110-aoncrg @ {:#x}", self.base);
        Ok(())
    }

    fn display_name(&self) -> &str {
        "Starfive JH7110 AONCRG"
    }
}

impl ClockController for Aoncrg {
    fn enable_clock(&self, clock: Option<u32>) -> Result<(), Error> {
        let index = clock.ok_or(Error::InvalidArgument)?;
        self.clk_enable_int(index)
    }

    fn disable_clock(&self, clock: Option<u32>) -> Result<(), Error> {
        log::warn!("jh7110: disable clock {clock:?}");
        Ok(())
    }
}

impl ResetController for Aoncrg {
    fn assert_reset(&self, reset: Option<u32>) -> Result<(), Error> {
        let reset = reset.ok_or(Error::InvalidArgument)?;
        let mapping = self.ensure_init()?;
        let regs = mapping.lock();
        let val = regs.RESETS.get();
        regs.RESETS.set(val | (1 << reset));
        Ok(())
    }

    fn deassert_reset(&self, reset: Option<u32>) -> Result<(), Error> {
        let reset = reset.ok_or(Error::InvalidArgument)?;
        let mapping = self.ensure_init()?;
        let regs = mapping.lock();
        let val = regs.RESETS.get();
        regs.RESETS.set(val & !(1 << reset));
        Ok(())
    }
}

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

impl DeviceTreeResetController for Aoncrg {
    fn map_reset(self: Arc<Self>, property: &TProp, offset: usize) -> Option<(ResetHandle, usize)> {
        let cell = property.read_cell(offset, 1)? as u32;
        Some((
            ResetHandle {
                parent: self,
                reset: Some(cell),
            },
            1,
        ))
    }
}

device_tree_driver! {
    compatible: ["starfive,jh7110-aoncrg"],
    driver: {
        fn probe(&self, node: &Arc<Node>, context: &mut ProbeContext) -> Option<Arc<dyn Device>> {
            let base = node.map_base(context, 0)?;

            let clock_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 = Arc::new(Aoncrg {
                base,
                clock_parents,
                mapping: OneTimeInit::new()
            });
            node.make_reset_controller(aoncrg.clone());
            node.make_clock_controller(aoncrg.clone());
            Some(aoncrg)
        }
    }
}