use libk_mm_interface::table::{EntryLevel, EntryLevelExt};

use crate::mem::table::L3;

pub fn tlb_flush_global_full() {
    tlb_flush_full();
    // TODO send TLB shootdown IPI to other harts
}

pub fn tlb_flush_global_va(va: usize) {
    tlb_flush_va(va);
    // TODO send TLB shootdown IPI to other harts
}

pub fn tlb_flush_range_va(start: usize, size: usize) {
    let end = (start + size).page_align_up::<L3>();
    let start = start.page_align_down::<L3>();

    for page in (start..end).step_by(L3::SIZE) {
        tlb_flush_va(page);
    }
}

pub fn tlb_flush_range_va_asid(asid: usize, start: usize, size: usize) {
    let end = (start + size).page_align_up::<L3>();
    let start = start.page_align_down::<L3>();

    for page in (start..end).step_by(L3::SIZE) {
        tlb_flush_va_asid(page, asid);
    }
}

#[inline]
pub fn tlb_flush_full() {
    unsafe { core::arch::asm!("sfence.vma") };
}

#[inline]
pub fn tlb_flush_va(va: usize) {
    unsafe { core::arch::asm!("sfence.vma {0}, zero", in(reg) va) };
}

#[inline]
pub fn tlb_flush_asid(asid: usize) {
    unsafe { core::arch::asm!("sfence.vma zero, {0}", in(reg) asid) };
}

#[inline]
pub fn tlb_flush_va_asid(va: usize, asid: usize) {
    unsafe { core::arch::asm!("sfence.vma {0}, {1}", in(reg) va, in(reg) asid) };
}
