use alloc::sync::Arc;
use device_api::{
    clock::{ClockController, ClockHandle, Hertz},
    device::{Device, DeviceInitContext},
};
use device_tree::{
    DeviceTreePropertyRead, TProp,
    driver::{DeviceTreeClockController, Node, ProbeContext, device_tree_driver},
};
use libk::error::Error;
use libk_mm::device::DeviceMemoryIo;
use libk_util::{OneTimeInit, sync::IrqSafeSpinlock};
use tock_registers::{
    interfaces::{Readable, Writeable},
    register_structs,
    registers::ReadWrite,
};

use crate::pll::{PllCfg, WrpllData};

const CLK_COREPLL: u32 = 0;
const CLK_DDRPLL: u32 = 1;
const CLK_GEMGXLPLL: u32 = 2;
const CLK_TLCLK: u32 = 3;

register_structs! {
    Regs {
        (0x000 => hfxosccfg: ReadWrite<u32>),
        (0x004 => corepllcfg0: ReadWrite<u32, PllCfg::Register>),
        (0x008 => _0),
        (0x00C => ddrpllcfg0: ReadWrite<u32>),
        (0x010 => ddrpllcfg1: ReadWrite<u32>),
        (0x014 => _1),
        (0x01C => gemgxlpllcfg0: ReadWrite<u32>),
        (0x020 => gemgxlpllcfg1: ReadWrite<u32>),
        (0x024 => coreclksel: ReadWrite<u32>),
        (0x028 => deviceresetreg: ReadWrite<u32>),
        (0x02C => _2),
        (0x100 => @END),
    }
}

struct Prci {
    clk_osc: ClockHandle,
    clk_rtc: ClockHandle,
    regs: IrqSafeSpinlock<DeviceMemoryIo<'static, Regs>>,
    init: OneTimeInit<()>,
}

impl Regs {
    fn read_core_pll(&self) -> WrpllData {
        let pll = self.corepllcfg0.extract();
        let divr = pll.read(PllCfg::PLLR);
        let divf = pll.read(PllCfg::PLLF);
        let divq = pll.read(PllCfg::PLLQ);
        WrpllData {
            divq,
            divr,
            divf,
            int_feedback: true,
        }
    }
}

impl Prci {
    fn ensure_init(&self) {
        // U-Boot should configure coreclk to 1GHz, but when booting
        // in QEMU, there is not U-Boot stage, and some peripherals
        // like PWM expect the clock to be configured this way. So
        // this needs to be done manually here
        self.init.or_init_with(|| {
            let regs = self.regs.lock();
            if regs.coreclksel.get() != 0 {
                // coreclk configured to hfclk
                // configure corepll to output 1GHz and switch to it
                // divf = 150
                // divr = 4
                // divq = 1
                regs.corepllcfg0.write(
                    PllCfg::PLLF.val(150)
                        + PllCfg::PLLR.val(4)
                        + PllCfg::PLLQ.val(1)
                        + PllCfg::PLLFSEBYPASS::SET,
                );
                while !regs.corepllcfg0.matches_all(PllCfg::PLLLOCK::SET) {
                    core::hint::spin_loop();
                }
                regs.coreclksel.set(0);
            }
        });
    }
}

impl Device for Prci {
    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        self.ensure_init();
        if let Ok(coreclk) = self.clock_rate(Some(CLK_COREPLL)) {
            log::info!("coreclk = {coreclk}");
        }
        Ok(())
    }

    fn display_name(&self) -> &str {
        "sifive,fu540-c000-prci"
    }
}

impl ClockController for Prci {
    fn clock_rate(&self, clock: Option<u32>) -> Result<Hertz, Error> {
        self.ensure_init();
        let _ = &self.clk_rtc;
        match clock {
            Some(CLK_COREPLL) => {
                let regs = self.regs.lock();
                let hfclk = self.clk_osc.rate()?;
                // == 0 -> use corepll
                // == 1 -> use hfclk
                let coreclksel = regs.coreclksel.get() != 0;
                let coreclk = match coreclksel {
                    true => hfclk,
                    false => {
                        let pll = regs.read_core_pll();
                        pll.output_rate(hfclk)
                    }
                };
                Ok(coreclk)
            }
            Some(CLK_TLCLK) => {
                let coreclk = self.clock_rate(Some(CLK_COREPLL))?;
                Ok(coreclk / 2)
            }
            Some(_) => todo!(),
            None => Err(Error::InvalidArgument),
        }
    }

    fn set_clock_rate(&self, _clock: Option<u32>, _rate: Hertz) -> Result<Hertz, Error> {
        Err(Error::NotImplemented)
    }

    fn enable_clock(&self, clock: Option<u32>) -> Result<(), Error> {
        self.ensure_init();
        match clock {
            Some(CLK_GEMGXLPLL) | Some(CLK_DDRPLL) | Some(CLK_COREPLL) => todo!(),
            Some(CLK_TLCLK) => Ok(()),
            _ => Err(Error::InvalidArgument),
        }
    }

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

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

device_tree_driver! {
    compatible: ["sifive,fu540-c000-prci"],
    driver: {
        fn probe(&self, node: &Arc<Node>, context: &mut ProbeContext) -> Option<Arc<dyn Device>> {
            let base = node.map_base(context, 0)?;
            let clk_osc = node.clock(0)?;
            let clk_rtc = node.clock(1)?;

            let regs = unsafe { DeviceMemoryIo::map(base, Default::default()) }.ok()?;

            let prci = Arc::new(Prci {
                regs: IrqSafeSpinlock::new(regs),
                clk_osc,
                clk_rtc,
                init: OneTimeInit::new()
            });

            node.make_clock_controller(prci.clone());

            Some(prci)
        }
    }
}
