//! Starfive JH7110 System Control Registers

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::DeviceMemoryIoMut};
use libk_util::{sync::IrqSafeSpinlock, OneTimeInit};

const PARENT_CLK_OSC: usize = 0;

struct Syscrg {
    base: PhysicalAddress,
    parents: [ClockHandle; 1],
    mapping: OneTimeInit<IrqSafeSpinlock<DeviceMemoryIoMut<'static, [u32]>>>,
}

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

#[derive(Debug, Clone, Copy)]
enum ClockRegister {
    // Clock gate
    Gate(u32, ClockParent),
    // Clock inverter
    Inv(u32, ClockParent),
    // Clock divider (TODO)
    Div(ClockParent),
    Unimp,
}

impl Syscrg {
    fn ensure_init(&self) -> Result<&IrqSafeSpinlock<DeviceMemoryIoMut<'static, [u32]>>, Error> {
        self.mapping.or_try_init_with(move || {
            unsafe { DeviceMemoryIoMut::map_slice(self.base, 256, Default::default()) }
                .map(IrqSafeSpinlock::new)
        })
    }

    fn clk_enable_int(&self, index: u32, depth: 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::Gate(bit, parent) => {
                log::info!("jh7110-syscrg: enable {name:?} @ {index} /{depth}");
                self.clk_enable_parent(parent, depth)?;
                let mut lock = regs.lock();
                lock[index as usize] |= 1 << bit;
                Ok(())
            }
            ClockRegister::Inv(bit, parent) => {
                log::info!("jh7110-syscrg: enable clk inv {name:?} @ {index}");
                self.clk_enable_parent(parent, depth)?;
                let mut lock = regs.lock();
                lock[index as usize] |= 1 << bit;
                Ok(())
            }
            ClockRegister::Div(parent) => self.clk_enable_parent(parent, depth),
            ClockRegister::Unimp => {
                log::warn!("jh7110-syscrg: clock not implemented: {name:?}");
                Err(Error::NotImplemented)
            }
        }
    }

    fn clk_enable_parent(&self, parent: ClockParent, depth: u32) -> Result<(), Error> {
        match parent {
            ClockParent::Int(index) => self.clk_enable_int(index, depth + 1),
            ClockParent::Ext(clock) => self.parents.get(clock).ok_or(Error::DoesNotExist)?.enable(),
            ClockParent::None => Ok(()),
        }
    }

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

            let mut t = [("", Unimp); 256];

            t[0x01C / 4] = ("clk_axi_cfg0", Div(None));
            t[0x020 / 4] = ("clk_stg_axiahb", Div(Int(0x01C / 4)));
            t[0x024 / 4] = ("clk_ahb0", Gate(31, Int(0x020 / 4)));

            t[0x184 / 4] = ("clk_gmac5_axi64_ahb", Gate(31, Int(0x024 / 4)));
            t[0x188 / 4] = ("clk_gmac5_axi64_axi", Gate(31, Int(0x020 / 4)));
            t[0x18C / 4] = ("clk_gmac_source", Div(None));
            t[0x190 / 4] = ("clk_gmac1_gtx", Gate(31, None));
            t[0x194 / 4] = ("clk_gmac1_rmii_rtx", Unimp);
            t[0x198 / 4] = ("clk_gmac5_axi64_ptp", Gate(31, Int(0x18C / 4)));
            t[0x19C / 4] = ("clk_gmac5_axi64_rx", Unimp);
            t[0x1A0 / 4] = ("clk_gmac5_axi64_rx_inv", Inv(30, Int(0x19C / 4)));
            t[0x1A4 / 4] = ("clk_gmac5_axi64_tx", Gate(31, None));
            t[0x1A8 / 4] = ("clk_gmac5_axi64_tx_inv", Inv(30, Int(0x1A4 / 4)));
            t[0x1AC / 4] = ("clk_gmac1_gtxc", Gate(31, Int(0x190 / 4)));
            t[0x1B0 / 4] = ("clk_gmac0_gtx", Gate(31, None));
            t[0x1B4 / 4] = ("clk_gmac0_ptp", Gate(31, Int(0x18C / 4)));
            t[0x1B8 / 4] = ("clk_gmac_phy", Gate(31, None));
            t[0x1BC / 4] = ("clk_gmac0_gtxc", Gate(31, Int(0x1B0 / 4)));

            t[0x244 / 4] = ("clk_u0_uart_apb", Gate(31, None));
            t[0x248 / 4] = ("clk_u0_uart_core", Gate(31, Ext(PARENT_CLK_OSC)));

            t
        };

        let index = index as usize;
        CLOCKS.get(index).copied()
    }
}

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

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

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

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

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

        match reg {
            ClockRegister::Gate(_, parent) => match parent {
                ClockParent::Ext(n) => self.parents[n].rate(),
                ClockParent::Int(_) => {
                    log::warn!("jh7110-syscrg: todo internal clock parents: {:?}", name);
                    Err(Error::NotImplemented)
                }
                ClockParent::None => {
                    log::warn!("jh7110-syscrg: clock parent not specified {:?}", name);
                    Err(Error::NotImplemented)
                }
            },
            ClockRegister::Unimp | ClockRegister::Div(_) | ClockRegister::Inv(_, _) => {
                log::warn!("jh7110-syscrg: unimplemented clock {:?}", name);
                Err(Error::NotImplemented)
            }
        }
    }

    fn set_clock_rate(&self, clock: Option<u32>, rate: u64) -> Result<u64, Error> {
        log::warn!("TODO: jh7110-syscrg: set rate {clock:?} -> {rate}");
        Ok(rate)
    }
}

impl Syscrg {
    fn rst_trigger(&self, reset: u32, assert: bool) -> Result<(), Error> {
        let reg = (190 + reset / 32) as usize;
        let bit = reset % 32;
        let regs = self.ensure_init()?;
        let mut regs = regs.lock();

        let expect = if assert {
            regs[reg] |= 1 << bit;
            1 << bit
        } else {
            regs[reg] &= !(1 << bit);
            0
        };

        let mut timeout = 1 << 20;
        while timeout > 0 && regs[reg] & (1 << bit) != expect {
            core::hint::spin_loop();
            timeout -= 1;
        }

        if timeout > 0 {
            Ok(())
        } else {
            log::warn!("jh7110-syscrg: reset timeout {reset}");
            Err(Error::TimedOut)
        }
    }
}

impl ResetController for Syscrg {
    fn assert_reset(&self, reset: Option<u32>) -> Result<(), Error> {
        let reset = reset.ok_or(Error::InvalidArgument)?;
        self.rst_trigger(reset, true)
    }

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

impl DeviceTreeClockController for Syscrg {
    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 DeviceTreeResetController for Syscrg {
    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,
        ))
    }
}

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

            let syscrg = Arc::new(Syscrg {
                base,
                parents: [osc],
                mapping: OneTimeInit::new(),
            });
            node.make_reset_controller(syscrg.clone());
            node.make_clock_controller(syscrg.clone());
            Some(syscrg)
        }
    }
}