//! The implementation of the time management trait for aarch64.

use core::arch::asm;
use core::time::Duration;

use crate::traits::time::TimeManager;
use crate::println;

const NS_PER_S: u64 = 1_000_000_000;
const CNTP_CTL_EL0_ENABLE_MASK: u64 = 0b001;
const CNTP_CTL_EL0_IMASK_MASK: u64 = 0b010;
const CNTP_CTL_EL0_ISTATUS_MASK: u64 = 0b100;

/// AARCH64 Timer.
struct Timer;

/// The global timer.
static TIME_MANAGER: Timer = Timer;

/// Return a reference to the global time manager.
pub fn time_manager() -> &'static impl TimeManager {
    &TIME_MANAGER
}

impl TimeManager for Timer {
    fn frequency(&self) -> u64 {
        let freq: u64;
        unsafe {
            asm!("mrs {0:x}, CNTFRQ_EL0", out(reg) freq, options(nomem, nostack));
        }
        freq
    }

    fn uptime(&self) -> Duration {
        let ticks: u64;
        unsafe {
            // Protect against out of order execution
            // https://developer.arm.com/documentation/dui0802/b/CIHGHHIE
            // In effect this flush the processor pipeline
            asm!("ISB", options(nostack));
            asm!("mrs {0:x}, CNTPCT_EL0", out(reg) ticks, options(nomem, nostack));
        }
        Duration::from_nanos((ticks*NS_PER_S)/self.frequency())
    }

    fn sleep(&self, duration: Duration) {
        if duration.as_nanos() == 0 {
            return;
        }

        let tmp = match self.frequency().checked_mul(duration.as_nanos() as u64) {
            None => {
                println!("Cannot sleep this long, skipping");
                return;
            }
            Some(tmp) => tmp,
        };
        let ticks = tmp / NS_PER_S;

        if ticks > u32::max_value().into() {
            println!("Cannot wait more than 2^32-1 ticks on aarch64, skipping");
            return;
        }
        if ticks == 0 {
            println!("Duration is smaller than a clock tick, skipping.");
            return;
        }

        // Enable the timer and mask (disable) the interrupt 
        let mut ctl: u64 = CNTP_CTL_EL0_ENABLE_MASK + CNTP_CTL_EL0_IMASK_MASK;
        unsafe {
            // set the register values for the timer
            asm!("msr CNTP_TVAL_EL0, {0:x}", in(reg) ticks, options(nomem, nostack));
            asm!("msr CNTP_CTL_EL0, {0:x}", in(reg) ctl, options(nomem, nostack));
        }
        // Loop waiting for the timer to run out
        while (ctl & CNTP_CTL_EL0_ISTATUS_MASK) == 0 {
            unsafe {
                asm!("mrs {0:x}, CNTP_CTL_EL0", out(reg) ctl, options(nomem, nostack));
            }
        }
        // Disable the counter
        ctl &= !CNTP_CTL_EL0_ENABLE_MASK;
        unsafe {
            asm!("msr CNTP_CTL_EL0, {0:x}", in(reg) ctl, options(nomem, nostack));
        }
    }
}
