//! RISC-V boot code

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

use elf::{abi, relocation::Elf64_Rela};
use kernel_arch::{Architecture, ArchitectureImpl};
use kernel_arch_riscv64::{mem, BOOT_HART_ID, CPU_COUNT};
use libk::{
    debug,
    fs::{devfs, sysfs},
    task::runtime,
};
use libk_mm::{
    address::{PhysicalAddress, Virtualize},
    PageBox,
};

use crate::{
    arch::{riscv64::smp::SecondaryContext, PLATFORM},
    kernel_main, kernel_secondary_main,
    mem::KERNEL_VIRT_OFFSET,
};

const BOOT_STACK_SIZE: usize = 65536;

static mut KERNEL_LOAD_BASE: u64 = 0;
static mut DTB_ADDRESS: PhysicalAddress = PhysicalAddress::ZERO;
static mut BOOT_STACK: BootStack<BOOT_STACK_SIZE> = BootStack::zeroed();

#[repr(C, align(0x10))]
struct BootStack<const N: usize>([u8; N]);

impl<const N: usize> BootStack<N> {
    pub const fn zeroed() -> Self {
        Self([0; N])
    }
}

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_RISCV_RELATIVE => rela.r_addend + image_base,
            _ => ArchitectureImpl::halt(),
        };

        slot.write_volatile(value);
    }
}

unsafe fn long_jump(pc: usize, sp: usize, a0: usize) -> ! {
    core::arch::asm!("mv sp, {sp}; jr {pc}", in("a0") a0, 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);

    debug::init_logger();
    super::debug::register_sbi_debug();

    log::info!("Starting riscv64 upper half");

    if DTB_ADDRESS.is_zero() {
        log::error!("No device tree provided");
        // No DTB provided
        ArchitectureImpl::halt();
    }

    if let Err(error) = PLATFORM.init_memory_management(DTB_ADDRESS) {
        log::error!("Failed to initialize memory management: {error:?}");
        ArchitectureImpl::halt();
    }

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

    runtime::init_task_queue();

    if let Err(error) = PLATFORM.init_platform(BOOT_HART_ID as u32, 0, true) {
        log::error!("Failed to initialize the platform: {error:?}");
        ArchitectureImpl::halt();
    }

    kernel_main()
}

unsafe extern "C" fn bsp_smode_entry() -> ! {
    mem::enable_mmu();
    let pc = bsp_entry_upper as usize + KERNEL_VIRT_OFFSET;
    let sp = (&raw const BOOT_STACK).addr() + BOOT_STACK_SIZE + KERNEL_VIRT_OFFSET;
    long_jump(pc, sp, 0)
}

unsafe extern "C" fn ap_smode_upper(context: PhysicalAddress) -> ! {
    let hart_id = {
        let context = PageBox::<SecondaryContext>::from_physical_raw(context);
        context.hart_id
    };
    let queue_index = CPU_COUNT.fetch_add(1, Ordering::Acquire);

    if let Err(error) = PLATFORM.init_platform(hart_id as u32, queue_index, false) {
        log::error!("Secondary hart init error: {error:?}");
        ArchitectureImpl::halt();
    }

    kernel_secondary_main()
}

unsafe extern "C" fn ap_smode_entry(context: PhysicalAddress, sp: PhysicalAddress) -> ! {
    compiler_fence(Ordering::SeqCst);
    mem::enable_mmu();
    compiler_fence(Ordering::SeqCst);
    let pc = ap_smode_upper as usize + KERNEL_VIRT_OFFSET;
    let sp = sp.virtualize();
    long_jump(pc, sp, context.into_usize())
}

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

    dtb_address = sym DTB_ADDRESS,
    kernel_load_base = sym KERNEL_LOAD_BASE,
    boot_hart_id = sym BOOT_HART_ID,

    boot_stack_bottom = sym BOOT_STACK,
    boot_stack_size = const BOOT_STACK_SIZE,
    bsp_smode_entry = sym bsp_smode_entry,

    ap_smode_entry = sym ap_smode_entry,

    relocate_kernel = sym relocate_kernel,
);
