//! RISC-V system timer driver
use core::sync::atomic::{AtomicU64, Ordering};

use abi::time::NANOSECONDS_IN_SECOND;
use kernel_arch::task::Scheduler;
use kernel_arch_riscv64::{intrinsics, registers::SIE, sbi};
use libk::{arch::Cpu, task::runtime, time};
use tock_registers::interfaces::ReadWriteable;

/// Timer value at last tick
pub static LAST_TICK: AtomicU64 = AtomicU64::new(0);
/// Timer frequency
pub static FREQUENCY: AtomicU64 = AtomicU64::new(0);
// 1kHz
const TICK_RATE: u64 = 1000;

// TODO use stimecmp instead of sbi calls if sstc extension is available

/// Handles incoming timer interrupts
pub fn handle_interrupt() {
    let frequency = FREQUENCY.load(Ordering::Acquire);

    // TODO merge this code with other system timer implementations
    let now = intrinsics::rdtime();
    sbi::sbi_set_timer(now.wrapping_add(frequency / TICK_RATE));

    if Cpu::local().is_bootstrap() {
        let last = LAST_TICK.swap(now, Ordering::Release);
        if frequency != 0 {
            if let Some(delta) = now.checked_sub(last) {
                // Only update time from local CPU
                let dt = delta * NANOSECONDS_IN_SECOND / frequency;
                time::add_nanoseconds(dt);
            }
        }
        LAST_TICK.store(now, Ordering::Release);

        runtime::tick();
    }

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

/// Sets up a timer for the calling HART
pub fn init_hart(is_bsp: bool) {
    let frequency = FREQUENCY.load(Ordering::Acquire);

    if frequency != 0 {
        SIE.modify(SIE::STIE::SET);
        let now = intrinsics::rdtime();
        if is_bsp {
            LAST_TICK.store(now, Ordering::Release);
        }
        sbi::sbi_set_timer(now.wrapping_add(frequency / TICK_RATE));
    } else {
        SIE.modify(SIE::STIE::CLEAR);
        sbi::sbi_set_timer(u64::MAX);
    }
}
