use core::arch::asm;

use aarch64_cpu::{asm::barrier, registers::*};
use aarch64_cpu_ext::asm::tlb::*;
use page_table_generic::VirtAddr;
use pie_boot_macros::start_code;

use crate::mem::PageTable;

#[start_code]
pub fn switch_to_elx(bootargs: usize) {
    SPSel.write(SPSel::SP::ELx);
    SP_EL0.set(0);
    let current_el = CurrentEL.read(CurrentEL::EL);
    let ret = sym_lma!(super::_start_secondary);
    if current_el == 3 {
        // Set EL2 to 64bit and enable the HVC instruction.
        SCR_EL3.write(
            SCR_EL3::NS::NonSecure + SCR_EL3::HCE::HvcEnabled + SCR_EL3::RW::NextELIsAarch64,
        );
        // Set the return address and exception level.
        SPSR_EL3.write(
            SPSR_EL3::M::EL1h
                + SPSR_EL3::D::Masked
                + SPSR_EL3::A::Masked
                + SPSR_EL3::I::Masked
                + SPSR_EL3::F::Masked,
        );

        ELR_EL3.set(ret as _);
        barrier::isb(barrier::SY);
        unsafe {
            asm!(
                "
                    mov x0, {}
                    eret
                    ",
                in(reg) bootargs,
                options(nostack, noreturn),
            );
        }
    }
}
#[inline(always)]
pub(crate) fn flush_tlb(vaddr: Option<VirtAddr>) {
    match vaddr {
        Some(addr) => {
            tlbi(VAE2IS::new(0, addr.raw()));
        }
        None => {
            tlbi(ALLE2);
        }
    }
    barrier::dsb(barrier::SY);
    barrier::isb(barrier::SY);
}

pub fn get_kernal_table() -> PageTable {
    let val = TTBR0_EL2.extract();
    PageTable {
        id: 0,
        addr: val.get() as _,
    }
}

pub fn set_kernal_table(tb: PageTable) {
    TTBR0_EL2.set(tb.addr as u64);
    tlbi(ALLE2);
    barrier::dsb(barrier::SY);
    barrier::isb(barrier::SY);
}
