//! AArch64 boot and entry implementation

use core::{
    arch::global_asm,
    sync::atomic::{self, AtomicUsize, Ordering},
};

use aarch64_cpu::{
    asm::barrier,
    registers::{
        CNTHCTL_EL2, CNTHP_CTL_EL2, CNTKCTL_EL1, CNTVOFF_EL2, CPACR_EL1, ELR_EL2, HCR_EL2,
        SPSR_EL2, SP_EL1, TTBR0_EL1,
    },
};
use elf::{abi, relocation::Elf64_Rela};
use kernel_arch::{Architecture, ArchitectureImpl, KERNEL_VIRT_OFFSET};
use kernel_arch_aarch64::{
    mem::{self, table::L3},
    CPU_COUNT,
};
use libk::{
    fs::{devfs, sysfs},
    task::runtime,
};
use libk_mm::{
    address::{PhysicalAddress, Virtualize},
    phys,
    table::EntryLevel,
};
use tock_registers::interfaces::{ReadWriteable, Writeable};

use crate::{
    arch::{
        aarch64::{exception, BootStack, BOOT_STACK_SIZE},
        PLATFORM,
    },
    kernel_main, kernel_secondary_main,
};

static BSP_STACK: BootStack<BOOT_STACK_SIZE> = BootStack::zeroed();
pub(crate) static SPIN_TABLE_STACK: AtomicUsize = AtomicUsize::new(0);
static mut KERNEL_LOAD_BASE: u64 = 0;
static mut DTB_PHYSICAL_ADDRESS: PhysicalAddress = PhysicalAddress::ZERO;

unsafe extern "C" fn relocate_kernel(image_base: i64, rela_start: usize, rela_end: usize) {
    let rela_count = (rela_end - rela_start) / size_of::<Elf64_Rela>();
    let rela_ptr: *const Elf64_Rela = core::ptr::with_exposed_provenance(rela_start);
    let rela_table = core::slice::from_raw_parts(rela_ptr, rela_count);

    for rela in rela_table {
        let slot: *mut i64 =
            core::ptr::with_exposed_provenance_mut(rela.r_offset as usize + image_base as usize);

        let value = match rela.r_info as u32 {
            abi::R_AARCH64_RELATIVE => rela.r_addend + image_base,
            _ => ArchitectureImpl::halt(),
        };

        slot.write_volatile(value);
    }
}

fn setup_cpu_common() {
    CPACR_EL1.modify(CPACR_EL1::FPEN::TrapNothing);
}

unsafe fn leave_el2(pc: usize, sp: usize, x0: usize) -> ! {
    CNTHCTL_EL2.write(CNTHCTL_EL2::EL1PCEN::SET + CNTHCTL_EL2::EL1PCTEN::SET);
    CNTHP_CTL_EL2.write(CNTHP_CTL_EL2::ENABLE::SET);
    CNTKCTL_EL1.set(0);
    CNTVOFF_EL2.set(0);

    HCR_EL2.write(HCR_EL2::RW::EL1IsAarch64 + HCR_EL2::SWIO::SET);

    SPSR_EL2.write(
        SPSR_EL2::D::SET
            + SPSR_EL2::A::SET
            + SPSR_EL2::I::SET
            + SPSR_EL2::F::SET
            + SPSR_EL2::M::EL1h,
    );

    ELR_EL2.set(pc as u64);
    SP_EL1.set(sp as u64);
    core::arch::asm!("eret", in("x0") x0, options(noreturn));
}

#[inline]
unsafe fn long_jump(pc: usize, sp: usize) -> ! {
    core::arch::asm!("mov sp, {sp}; br {pc}", sp = in(reg) sp, pc = in(reg) pc, options(noreturn));
}

unsafe extern "C" fn bsp_entry_upper() -> ! {
    // Relocate the kernel again
    extern "C" {
        static __rela_start: u8;
        static __rela_end: u8;
    }

    atomic::fence(Ordering::SeqCst);
    let rela_start = (&raw const __rela_start).addr();
    let rela_end = (&raw const __rela_end).addr();
    let image_base = (KERNEL_VIRT_OFFSET as u64 + KERNEL_LOAD_BASE) as i64;
    relocate_kernel(image_base, rela_start, rela_end);
    atomic::fence(Ordering::SeqCst);

    // Setup memory management (using DTB from lower half)
    if let Err(error) = PLATFORM.init_memory_management(DTB_PHYSICAL_ADDRESS) {
        let _ = error;
        ArchitectureImpl::halt();
    }

    // Remove the "lower-half" mapping, no longer needed
    TTBR0_EL1.set(0);

    mem::tlb_flush_all();

    barrier::dsb(barrier::ISHST);
    barrier::isb(barrier::SY);

    exception::init_exceptions();

    sysfs::init();
    devfs::init();

    runtime::init_task_queue();

    // Initialize the BSP CPU + the devices
    if PLATFORM.init_platform(true).is_err() {
        ArchitectureImpl::halt();
    }

    kernel_main()
}

unsafe extern "C" fn ap_entry_upper() -> ! {
    // Flush lower half
    TTBR0_EL1.set(0);

    mem::tlb_flush_all();

    let cpu_id = CPU_COUNT.fetch_add(1, Ordering::SeqCst);
    barrier::dsb(barrier::ISHST);
    barrier::isb(barrier::SY);
    aarch64_cpu::asm::sev();
    barrier::isb(barrier::SY);

    log::info!("cpu{} initializing", cpu_id);

    exception::init_exceptions();

    unsafe {
        PLATFORM
            .init_platform(false)
            .expect("Could not initialize the AP");
    }

    kernel_secondary_main()
}

unsafe extern "C" fn bsp_entry_el1() -> ! {
    setup_cpu_common();
    mem::init_lower(true);
    let pc = (bsp_entry_upper as usize) + KERNEL_VIRT_OFFSET;
    let sp = BSP_STACK.top_addr() + KERNEL_VIRT_OFFSET;
    long_jump(pc, sp)
}

unsafe extern "C" fn bsp_entry_el2() -> ! {
    leave_el2(bsp_entry_el1 as usize, BSP_STACK.top_addr(), 0)
}

unsafe extern "C" fn ap_entry_el1(sp: usize) -> ! {
    let _ = sp;
    const AP_STACK_PAGES: usize = 8;

    setup_cpu_common();
    mem::init_lower(false);

    let stack_pages = phys::alloc_pages_contiguous(AP_STACK_PAGES).unwrap();
    let stack_base = stack_pages.virtualize();
    let sp = stack_base + L3::SIZE * AP_STACK_PAGES;

    let pc = (ap_entry_upper as usize) + KERNEL_VIRT_OFFSET;
    long_jump(pc, sp)
}

unsafe extern "C" fn ap_entry_el2(sp: usize) -> ! {
    leave_el2(ap_entry_el1 as usize, sp, sp)
}

global_asm!(
    include_str!("entry.S"),

    kernel_load_base = sym KERNEL_LOAD_BASE,
    dtb_physical_address = sym DTB_PHYSICAL_ADDRESS,

    bsp_stack_size = const BOOT_STACK_SIZE,
    bsp_stack_bottom = sym BSP_STACK,
    bsp_entry_el1 = sym bsp_entry_el1,
    bsp_entry_el2 = sym bsp_entry_el2,

    ap_entry_el1 = sym ap_entry_el1,
    ap_entry_el2 = sym ap_entry_el2,

    relocate_kernel = sym relocate_kernel,
);
