use moa_mm::PTRS_PER_TABLE;

const MAX_TLBI_OPS: usize = PTRS_PER_TABLE;

#[inline(always)]
const fn tlbi_addr(addr: usize, asid: usize) -> usize {
    ((addr >> 12) & ((1 << 44) - 1)) | (asid << 48)
}

#[inline(always)]
pub(super) fn arch_local_flush_tlb_all() {
    unsafe {
        core::arch::aarch64::__dsb(core::arch::aarch64::NSHST);
        core::arch::asm!("tlbi vmalle1");
        core::arch::aarch64::__dsb(core::arch::aarch64::NSH);
        core::arch::aarch64::__isb(core::arch::aarch64::SY);
    }
}

#[inline(always)]
pub(super) fn arch_flush_tlb_space(asid: usize) {
    let arg = tlbi_addr(0, asid);
    unsafe {
        core::arch::aarch64::__dsb(core::arch::aarch64::ISHST);
        core::arch::asm!("tlbi aside1is, {0}", in(reg) arg);
        core::arch::aarch64::__dsb(core::arch::aarch64::ISH);
    }
}

pub(super) fn arch_flush_tlb_range(asid: usize, start: usize, end: usize, stride: usize) {
    if (end - start) / stride > MAX_TLBI_OPS {
        arch_flush_tlb_space(asid);
        return;
    }

    unsafe {
        core::arch::aarch64::__dsb(core::arch::aarch64::ISHST);
        let mut addr = start;
        while addr < end {
            let arg = tlbi_addr(addr, asid);
            core::arch::asm!("tlbi vae1is, {0}", in(reg) arg);
            addr += stride;
        }
        core::arch::aarch64::__dsb(core::arch::aarch64::ISH);
    }
}
