use core::sync::atomic::{self, Ordering};

use aarch64_cpu::{
    asm::barrier,
    registers::{PAR_EL1, SCTLR_EL1},
};
use libk_mm_interface::table::{EntryLevel, EntryLevelExt};
use tock_registers::interfaces::{ReadWriteable, Readable};

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

/// Enables data cache.
///
/// # Safety
///
/// Manipulates low-level machine state, use with care.
pub unsafe fn enable_dcache() {
    barrier::dsb(barrier::ISHST);
    barrier::isb(barrier::SY);

    SCTLR_EL1.modify(SCTLR_EL1::C::Cacheable);

    barrier::dsb(barrier::ISH);
    barrier::isb(barrier::SY);
}

/// Enables instruction cache.
///
/// # Safety
///
/// Manipulates low-level machine state, use with care.
pub unsafe fn enable_icache() {
    barrier::isb(barrier::SY);

    SCTLR_EL1.modify(SCTLR_EL1::I::Cacheable);

    barrier::dsb(barrier::ISH);
    barrier::isb(barrier::SY);
}

/// Disables instruction cache.
///
/// # Safety
///
/// Manipulates low-level machine state, use with care. Might break some instructions.
pub unsafe fn disable_icache() {
    barrier::isb(barrier::SY);

    ic_iallu();
    SCTLR_EL1.modify(SCTLR_EL1::I::NonCacheable);

    barrier::dsb(barrier::ISH);
    barrier::isb(barrier::SY);
}

#[inline]
pub fn tlb_flush_asid(asid: u8) {
    barrier::dsb(barrier::ISHST);
    let value = (asid as u64) << 48;
    unsafe {
        core::arch::asm!("tlbi aside1, {value}", value = in(reg) value);
    }
    barrier::dsb(barrier::ISH);
    barrier::isb(barrier::SY);
}

#[inline]
pub fn tlb_flush_all() {
    barrier::dsb(barrier::ISHST);
    unsafe {
        core::arch::asm!("tlbi vmalle1is");
    }
    barrier::dsb(barrier::ISH);
    barrier::isb(barrier::SY);
}

#[inline]
pub fn tlb_flush_vaae1(page: usize) {
    barrier::dsb(barrier::ISHST);
    let argument = page >> 12;
    unsafe {
        core::arch::asm!("tlbi vaae1, {argument}", argument = in(reg) argument);
    }
    barrier::dsb(barrier::ISH);
    barrier::isb(barrier::SY);
}

#[inline]
pub fn tlb_flush_range_va(base: usize, size: usize) {
    let end = (base + size).page_align_up::<L3>();
    let base = base.page_align_down::<L3>();
    let count = (end - base).page_count::<L3>();
    for i in 0..count {
        tlb_flush_vaae1(base + i * L3::SIZE);
    }
}

pub fn at_s1e0r(input: usize) -> Option<u64> {
    barrier::dsb(barrier::ISHST);
    unsafe {
        core::arch::asm!("at s1e0r, {address}", address = in(reg) input);
    }
    barrier::dsb(barrier::ISH);
    barrier::isb(barrier::SY);
    if PAR_EL1.matches_all(PAR_EL1::F::TranslationSuccessfull) {
        Some(PAR_EL1.read(PAR_EL1::PA))
    } else {
        None
    }
}

pub fn at_s1e1r(input: usize) -> Option<u64> {
    barrier::dsb(barrier::ISHST);
    unsafe {
        core::arch::asm!("at s1e1r, {address}", address = in(reg) input);
    }
    barrier::dsb(barrier::ISH);
    barrier::isb(barrier::SY);
    if PAR_EL1.matches_all(PAR_EL1::F::TranslationSuccessfull) {
        Some(PAR_EL1.read(PAR_EL1::PA))
    } else {
        None
    }
}

pub fn ic_iallu() {
    atomic::compiler_fence(Ordering::SeqCst);
    barrier::dsb(barrier::ISH);
    barrier::isb(barrier::SY);
    unsafe {
        core::arch::asm!("ic iallu");
    }
    barrier::isb(barrier::SY);
}

pub fn dc_cvac(input: usize) {
    barrier::dsb(barrier::ISHST);
    unsafe {
        core::arch::asm!("dc cvac, {address}", address = in(reg) input);
    }
}
