#![allow(static_mut_refs)]

use alloc::sync::Arc;

use crate::{
    irq::irqflags::irqs_disabled,
    println,
    sync::seqlock::SeqCount,
    time::{
        NSEC_PER_SEC,
        clocksource::{clocks_calc_max_nsecs, clocks_calc_mult_shift},
        hrtimer::{Hrtimer, HrtimerRestart, hrtimer_start},
    },
};

#[derive(Clone, Copy)]
struct ClockReadData {
    epoch_ns: u64,
    epoch_cyc: u64,
    sched_clock_mask: u64,
    read_sched_clock: fn() -> u64,
    mult: u32,
    shift: u32,
}

impl ClockReadData {
    const fn new() -> Self {
        Self {
            epoch_ns: 0,
            epoch_cyc: 0,
            sched_clock_mask: 0xff,
            read_sched_clock: dummy_read,
            mult: 1,
            shift: 1,
        }
    }
}

#[repr(align(64))]
struct ClockData {
    seq: SeqCount,
    read_data: [ClockReadData; 2],
    wrap_kt: u64,
    rate: usize,
    actual_read_sched_clock: fn() -> u64,
}

static mut CD: ClockData = ClockData {
    seq: SeqCount::new(),
    read_data: [ClockReadData::new(); 2],
    wrap_kt: 1,
    rate: 1,
    actual_read_sched_clock: dummy_read,
};

#[inline(always)]
fn dummy_read() -> u64 {
    0
}

#[inline(always)]
fn cyc_to_ns(cyc: u64, mult: u32, shift: u32) -> u64 {
    (cyc * u64::from(mult)) >> shift
}

fn update_clock_read_data(rd: &ClockReadData) {
    unsafe {
        CD.read_data[1] = *rd;

        CD.seq.write_seqcount_latch();
        CD.read_data[0] = *rd;
        CD.seq.write_seqcount_latch();
    }
}

fn update_sched_clock() {
    unsafe {
        let mut rd = CD.read_data[0];

        let cyc = (CD.actual_read_sched_clock)();
        let ns =
            rd.epoch_ns + cyc_to_ns((cyc - rd.epoch_cyc) & rd.sched_clock_mask, rd.mult, rd.shift);

        rd.epoch_ns = ns;
        rd.epoch_cyc = cyc;

        update_clock_read_data(&rd);
    }
}

pub(crate) fn sched_clock_register(read: fn() -> u64, mask: u64, rate: usize) {
    unsafe {
        assert!(irqs_disabled());

        let mut new_mult = 0;
        let mut new_shift = 0;
        clocks_calc_mult_shift(
            &mut new_mult,
            &mut new_shift,
            rate as u32,
            NSEC_PER_SEC as u32,
            3600,
        );

        CD.rate = rate;

        let mut max_cyc = 0;
        let wrap = clocks_calc_max_nsecs(new_mult, new_shift, 0, mask, &mut max_cyc);
        CD.wrap_kt = wrap;

        let mut rd = CD.read_data[0];

        let new_epoch = read();
        let cyc = (CD.actual_read_sched_clock)();
        let ns =
            rd.epoch_ns + cyc_to_ns((cyc - rd.epoch_cyc) & rd.sched_clock_mask, rd.mult, rd.shift);
        CD.actual_read_sched_clock = read;

        rd.read_sched_clock = read;
        rd.sched_clock_mask = mask;
        rd.mult = new_mult;
        rd.shift = new_shift;
        rd.epoch_cyc = new_epoch;
        rd.epoch_ns = ns;

        update_clock_read_data(&rd);

        let mut r_unit = ' ';
        let mut r = rate;
        if r >= 4000000 {
            r /= 1000000;
            r_unit = 'M';
        } else if r >= 1000 {
            r /= 1000;
            r_unit = 'k';
        }

        let res = cyc_to_ns(1, new_mult, new_shift);
        println!(
            "sched_clock: {} bits at {}{}Hz, resolution {}ns, wraps every {}ns",
            mask.count_ones(),
            r,
            r_unit,
            res,
            wrap
        );
    }
}

pub(crate) fn sched_clock_init() {
    update_sched_clock();

    let timer = Arc::new(Hrtimer::create(super::hrtimer::HrtimerMode::Rel, sched_clock_poll, None));
    unsafe {
        assert!(hrtimer_start(&timer, CD.wrap_kt));
    }
}

fn sched_clock_poll(timer: &Hrtimer) -> HrtimerRestart {
    update_sched_clock();
    unsafe {
        timer.forward_now(CD.wrap_kt);
    }
    HrtimerRestart::Restart
}

pub(crate) fn sched_clock() -> u64 {
    unsafe {
        loop {
            let seq = CD.seq.read_seqcount();
            let rd = CD.read_data[seq as usize & 1];
            let cyc = ((rd.read_sched_clock)() - rd.epoch_cyc) & rd.sched_clock_mask;
            let res = rd.epoch_ns + cyc_to_ns(cyc, rd.mult, rd.shift);
            if !CD.seq.read_seqcount_retry(seq) {
                return res;
            }
        }
    }
}
