//! AArch64 Generic Timer

use core::sync::atomic::{AtomicU64, Ordering};

use aarch64_cpu::registers::{CNTFRQ_EL0, CNTPCT_EL0, CNTP_CTL_EL0, CNTP_TVAL_EL0};
use abi::{error::Error, time::NANOSECONDS_IN_SECOND};
use alloc::sync::Arc;
use device_api::{
    device::{Device, DeviceInitContext},
    interrupt::{InterruptHandler, IrqHandle, IrqVector},
};
use device_tree::driver::{device_tree_driver, Node, ProbeContext};
use kernel_arch::task::Scheduler;
use libk::{arch::Cpu, task::runtime, time};
use tock_registers::interfaces::{ReadWriteable, Readable, Writeable};

/// ARM Generic Timer driver
pub struct ArmTimer {
    irq: IrqHandle,
}

/// ARM timer tick interval (in some time units?)
pub const TICK_INTERVAL: u64 = 250000;

static LAST_TICKS: AtomicU64 = AtomicU64::new(0);

impl InterruptHandler for ArmTimer {
    fn handle_irq(self: Arc<Self>, _vector: IrqVector) -> bool {
        let count = CNTPCT_EL0.get();

        CNTP_TVAL_EL0.set(TICK_INTERVAL);

        if Cpu::local().is_bootstrap() {
            let last = LAST_TICKS.swap(count, Ordering::Relaxed);
            let freq = CNTFRQ_EL0.get();
            if let Some(delta) = count.checked_sub(last) {
                // Only update time from local CPU
                let dt = delta * NANOSECONDS_IN_SECOND / freq;
                time::add_nanoseconds(dt);
            }
        }

        runtime::tick();

        unsafe {
            Cpu::local().scheduler().yield_cpu();
        }

        true
    }
}

impl Device for ArmTimer {
    fn display_name(&self) -> &str {
        "ARM Generic Timer"
    }

    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        CNTP_CTL_EL0.write(CNTP_CTL_EL0::ENABLE::SET + CNTP_CTL_EL0::IMASK::SET);
        Ok(())
    }

    unsafe fn init_irq(self: Arc<Self>) -> Result<(), Error> {
        log::info!("ARM Generic Timer frequency={}Hz", CNTFRQ_EL0.get());

        self.irq.register(self.clone())?;

        CNTP_CTL_EL0.modify(CNTP_CTL_EL0::IMASK::CLEAR);
        CNTP_TVAL_EL0.set(TICK_INTERVAL);

        self.irq.enable()?;

        Ok(())
    }
}

impl ArmTimer {
    /// Constructs an instance of ARM generic timer.
    ///
    /// # Safety
    ///
    /// The caller must ensure the function has not been called before.
    pub const unsafe fn new(irq: IrqHandle) -> Self {
        Self { irq }
    }
}

// TODO handle interrupt-parent
device_tree_driver! {
    compatible: ["arm,armv8-timer"],
    driver: {
        fn probe(&self, node: &Arc<Node>, _context: &mut ProbeContext) -> Option<Arc<dyn Device>> {
            let irq = node.interrupt(1)?;
            Some(Arc::new(unsafe { ArmTimer::new(irq) }))
        }
    }
}
