use alloc::sync::Arc;
use device_api::{
    clock::{ClockHandle, Hertz},
    device::{Device, DeviceInitContext},
    pwm::{PwmController, PwmControllerInfo},
};
use device_tree::driver::{Node, ProbeContext, device_tree_driver};
use libk::{device::manager::DEVICE_REGISTRY, error::Error};
use libk_mm::device::DeviceMemoryIo;
use libk_util::sync::IrqSafeSpinlock;
use tock_registers::{
    interfaces::{ReadWriteable, Readable, Writeable},
    register_bitfields, register_structs,
    registers::ReadWrite,
};
use yggdrasil_abi::time::NANOSECONDS_IN_SECOND;

const CHANNEL_COUNT: usize = 4;
const CMPWIDTH: usize = 16;
const SCALE_MAX: u32 = (1 << 4) - 1;
const MAX_PERIOD_CYCLES: u64 = (1u64 << (CMPWIDTH + 16)) - 1;
const MAX_COMPARE_VALUE: u32 = (1u32 << 16) - 1;

register_bitfields! {
    u32,
    PwmCfg [
        SCALE       OFFSET(0)   NUMBITS(4)  [],
        STICKY      OFFSET(8)   NUMBITS(1)  [],
        ZEROCMP     OFFSET(9)   NUMBITS(1)  [],
        DEGLITCH    OFFSET(10)  NUMBITS(1)  [],
        ENALWAYS    OFFSET(12)  NUMBITS(1)  [],
        ENONESHOT   OFFSET(13)  NUMBITS(1)  [],
        CMP_CENTER  OFFSET(16)  NUMBITS(4)  [],
        CMP_GANG    OFFSET(24)  NUMBITS(4)  [],
        CMP_IP      OFFSET(28)  NUMBITS(4)  [],
    ]
}

register_structs! {
    Regs {
        (0x00 => pwmcfg: ReadWrite<u32, PwmCfg::Register>),
        (0x04 => _0),
        (0x08 => pwmcount: ReadWrite<u32>),
        (0x0C => _1),
        (0x10 => pwms: ReadWrite<u32>),
        (0x14 => _2),
        (0x20 => pwmcmp: [ReadWrite<u32>; CHANNEL_COUNT]),
        (0x30 => @END),
    }
}

struct Pwm {
    name: &'static str,
    clock: ClockHandle,
    regs: IrqSafeSpinlock<DeviceMemoryIo<'static, Regs>>,
}

impl Regs {
    fn ns_to_compare_value(&self, input_clk: Hertz, ns: u64) -> Option<u32> {
        let scale = self.pwmcfg.read(PwmCfg::SCALE);
        let tick_duration_ns = NANOSECONDS_IN_SECOND / (input_clk.0 >> scale);
        let ticks = ns / tick_duration_ns;
        if ticks < MAX_COMPARE_VALUE as u64 {
            Some(ticks as u32)
        } else {
            None
        }
    }

    fn compare_value_to_ns(&self, input_clk: Hertz, cval: u32) -> u64 {
        let scale = self.pwmcfg.read(PwmCfg::SCALE);
        let tick_duration_ns = NANOSECONDS_IN_SECOND / (input_clk.0 >> scale);
        tick_duration_ns * cval as u64
    }
}

impl Device for Pwm {
    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        // pwm counter registers are (cmpwidth + 15) bits wide
        self.clock.enable()?;
        // Set all channels to max value and channel 0 sets the period
        {
            let regs = self.regs.lock();

            regs.pwmcfg
                .write(PwmCfg::ZEROCMP::SET + PwmCfg::ENALWAYS::SET);
            for channel in 0..CHANNEL_COUNT {
                regs.pwmcmp[channel].set(MAX_COMPARE_VALUE);
            }
        }

        DEVICE_REGISTRY.pwm.register_controller(self)?;

        Ok(())
    }

    fn display_name(&self) -> &str {
        self.name
    }
}

impl PwmController for Pwm {
    fn controller_info(&self) -> Result<PwmControllerInfo, Error> {
        Ok(PwmControllerInfo {
            channel_count: CHANNEL_COUNT,
            per_channel_period: false,
        })
    }

    fn channel_period(&self, _channel: u32) -> Result<u64, Error> {
        let input_clk = self.clock.rate()?;
        let regs = self.regs.lock();
        let compare_value = regs.pwmcmp[0].get();
        if compare_value == MAX_COMPARE_VALUE {
            return Ok(0);
        }
        let period_ns = regs.compare_value_to_ns(input_clk, compare_value);
        Ok(period_ns)
    }

    fn set_channel_period(&self, _channel: u32, period: u64) -> Result<u64, Error> {
        // Channel 0 sets the period
        let input_clk = self.clock.rate()?;
        let period_cycles = period * u64::from(input_clk) / NANOSECONDS_IN_SECOND;
        let mut prescaler = 0;
        if period_cycles > MAX_PERIOD_CYCLES {
            return Err(Error::InvalidArgument);
        }

        while (period_cycles >> prescaler) >= MAX_COMPARE_VALUE as u64 {
            prescaler += 1;
        }

        if prescaler > SCALE_MAX {
            log::warn!("Cannot fit {prescaler} into 4-bit prescaler");
            return Err(Error::InvalidArgument);
        }

        let period_ticks = period_cycles >> prescaler;

        let regs = self.regs.lock();

        regs.pwmcfg
            .modify(PwmCfg::SCALE.val(prescaler) + PwmCfg::ZEROCMP::SET);
        regs.pwmcmp[0].set(period_ticks as u32);

        let tick_duration_ns = NANOSECONDS_IN_SECOND / (input_clk.0 >> prescaler);
        let real_period_ns = tick_duration_ns * period_ticks;

        Ok(real_period_ns)
    }

    fn set_duty_cycle(&self, channel: u32, t: u64) -> Result<(), Error> {
        // Channel 0 sets the period
        if channel == 0 {
            return Ok(());
        }
        let input_clk = self.clock.rate()?;
        let regs = self.regs.lock();
        let cval = regs
            .ns_to_compare_value(input_clk, t)
            .ok_or(Error::InvalidArgument)?;
        // If the cval exceeds the "period" comparator, the PWM will never fire
        let cval = cval.min(regs.pwmcmp[0].get().saturating_sub(1));
        regs.pwmcmp[channel as usize].set(cval);
        Ok(())
    }

    fn duty_cycle(&self, channel: u32) -> Result<u64, Error> {
        if channel == 0 {
            return Ok(0);
        }
        let input_clk = self.clock.rate()?;
        let regs = self.regs.lock();
        Ok(regs.compare_value_to_ns(input_clk, regs.pwmcmp[channel as usize].get()))
    }

    fn stop_channel(&self, channel: u32) -> Result<(), Error> {
        if channel == 0 {
            return Ok(());
        }
        self.regs.lock().pwmcmp[channel as usize].set(MAX_COMPARE_VALUE);
        Ok(())
    }
}

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

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

            let pwm = Arc::new(Pwm {
                name,
                regs: IrqSafeSpinlock::new(regs),
                clock,
            });

            Some(pwm)
        }
    }
}
