#![no_std]
#![no_main]

mod board;
mod hw;
mod kernel_cell;
mod mmu;
mod platform;
mod trap;

use core::arch::global_asm;
use hw::uart::Uart;
use kernel_cell::KernelCell;

static KERNEL: KernelCell<kernel::Kernel<platform::RiscVPlatform>> = KernelCell::uninit();

pub fn kernel_ref() -> &'static mut kernel::Kernel<platform::RiscVPlatform> {
    KERNEL.get_mut()
}

global_asm!(
    r#"
.section .text.entry
.globl _start
_start:
    la sp, _stack_top
    call kmain
spin:
    wfi
    j spin
"#
);

unsafe extern "C" {
    static _bss_start: u8;
    static _bss_end: u8;
    static _heap_start: u8;
}

fn zero_bss() {
    unsafe {
        let start = &_bss_start as *const u8 as *mut u8;
        let end = &_bss_end as *const u8 as *mut u8;
        let len = end as usize - start as usize;
        core::ptr::write_bytes(start, 0, len);
    }
}

#[unsafe(no_mangle)]
extern "C" fn kmain() -> ! {
    zero_bss();

    let uart = Uart::new();
    uart.write_str("k-os: booted on RISC-V\n");

    let heap_start = unsafe { &_heap_start as *const u8 as usize };
    mmu::init_allocator(heap_start);
    let kernel_pt = mmu::init_kernel_page_table();
    mmu::activate(kernel_pt);
    uart.write_str("k-os: MMU enabled (Sv39)\n");

    trap::init();
    hw::timer::init();
    hw::plic::Plic::init();
    uart.write_str("k-os: traps and timer initialized\n");

    let riscv_platform = platform::RiscVPlatform::new(kernel_pt as *const mmu::PageTable as usize);
    let pool = domain::memory::page_pool::PagePool::new(
        types::address::PhysicalAddress::from_raw(board::USER_POOL_BASE),
        board::USER_POOL_PAGES,
    );
    KERNEL.init(kernel::Kernel::new(riscv_platform, pool));
    let kern = kernel_ref();

    let (_pid, _tid) = kern
        .boot(types::bounded::Priority::DEFAULT, types::time::TimeSlice::DEFAULT)
        .unwrap_or_else(|_| loop { unsafe { core::arch::asm!("wfi") }; });
    uart.write_str("k-os: kernel started\n");

    userspace_demo(&uart, kernel_pt);

    uart.write_str("k-os: demo complete\n");
    loop {
        unsafe { core::arch::asm!("wfi") };
    }
}

static USER_EXIT_PROGRAM: [u8; 12] = [
    0x93, 0x08, 0x00, 0x01, // addi a7, zero, 16 (SYS_EXIT)
    0x73, 0x00, 0x00, 0x00, // ecall
    0x6f, 0xf0, 0xdf, 0xff, // j -4
];

const USER_VA_RAW: usize = 0x10000;

fn userspace_demo(uart: &Uart, kernel_pt: &mmu::PageTable) {
    uart.write_str("k-os: --- userspace demo ---\n");

    let program_page = mmu::allocator::alloc_zeroed_page();
    unsafe {
        core::ptr::copy_nonoverlapping(
            USER_EXIT_PROGRAM.as_ptr(),
            program_page as *mut u8,
            USER_EXIT_PROGRAM.len(),
        );
    }

    let user_pt = mmu::create_user_page_table(
        kernel_pt,
        types::address::PhysicalAddress::from_raw(program_page),
        1,
        types::address::VirtualAddress::from_raw(USER_VA_RAW),
    );

    uart.write_str("k-os: user page table created, entering U-mode\n");

    mmu::activate(user_pt);

    let kernel_sp: usize;
    unsafe { core::arch::asm!("mv {0}, sp", out(reg) kernel_sp) };
    trap::set_kernel_stack(kernel_sp);

    let user_stack = USER_VA_RAW - 8;
    unsafe {
        core::arch::asm!(
            "csrc sstatus, {spp}",
            "csrs sstatus, {spie}",
            "csrw sepc, {entry}",
            "mv sp, {stack}",
            "sret",
            spp = in(reg) 1usize << 8,
            spie = in(reg) 1usize << 5,
            entry = in(reg) USER_VA_RAW,
            stack = in(reg) user_stack,
            options(noreturn),
        );
    }
}

#[panic_handler]
fn panic(info: &core::panic::PanicInfo) -> ! {
    let uart = Uart::new();
    uart.write_str("PANIC: ");
    if let Some(location) = info.location() {
        uart.write_str(location.file());
        uart.put_char(b':');
        uart.write_dec(location.line() as usize);
    }
    uart.put_char(b'\n');
    loop {
        unsafe { core::arch::asm!("wfi") };
    }
}