#![no_std]
#![no_main]

mod board;
mod boot;
mod boot_fs;
mod devices;
mod hw;
mod kernel_cell;
mod mmu;
mod platform;
mod smp;
mod trap;

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

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

fn kern_mut() -> &'static mut kernel::Kernel<platform::RiscVPlatform> {
    // SAFETY: boot runs with the kernel lock held and before any other
    // hart has started.
    unsafe { &mut *kernel_ptr() }
}

pub fn kernel_ptr() -> *mut kernel::Kernel<platform::RiscVPlatform> {
    KERNEL.as_ptr()
}

// Firmware enters `_start` (and SBI starts other harts at
// `_secondary_start`) at the physical load address with the MMU off and
// the hart id in a0, but the kernel is linked in the upper half. Using
// only PC-relative code, the stub turns on `boot_page_table`, which maps
// RAM twice (at its physical address, just so the next instruction still
// exists, and in the upper half) plus the MMIO window high, then jumps to
// its own high address and takes the hart's kernel stack. The kernel then
// switches every hart to the real table, which maps nothing low.
global_asm!(
    r#"
.section .text.entry
.globl _start
_start:
    lla t2, start_high_address
    j 1f
.globl _secondary_start
_secondary_start:
    lla t2, secondary_high_address
1:
    li t0, 8                # smp::MAX_HARTS
    bgeu a0, t0, spin
    mv s0, a0
    lla t0, boot_page_table
    srli t0, t0, 12
    li t1, 8 << 60
    or t0, t0, t1
    csrw satp, t0
    sfence.vma
    ld t0, 0(t2)
    jr t0

# sp = _stacks_start + (hart + 1) * smp::STACK_SLOT
start_high:
    call hart_stack
    mv a0, s0
    call kmain
    j spin
secondary_high:
    call hart_stack
    mv a0, s0
    call kmain_secondary
    j spin
hart_stack:
    lla sp, _stacks_start
    addi t0, s0, 1
    li t1, 0x81000
    mul t0, t0, t1
    add sp, sp, t0
    ret
spin:
    wfi
    j spin

.pushsection .rodata
.balign 8
start_high_address:
    .quad start_high
secondary_high_address:
    .quad secondary_high
.popsection

# Sv39 root of 1 GiB leaves; an entry is (PA >> 12) << 10 | flags, with
# flags V|R|W|A|D = 0xc7, plus X = 0xcf.
.pushsection .data
.balign 4096
boot_page_table:
    .quad 0
    .quad 0
    .quad 0x200000cf    # [2]   0x8000_0000 -> RAM, for the jump only
    .fill 253, 8, 0
    .quad 0x000000c7    # [256] PHYS_OFFSET + 0 -> MMIO
    .quad 0
    .quad 0x200000cf    # [258] PHYS_OFFSET + 0x8000_0000 -> RAM
    .fill 253, 8, 0
.popsection
"#
);

const SHARED_BUF_VA: usize = 0x30000;

#[unsafe(no_mangle)]
extern "C" fn kmain(hart: usize) -> ! {
    let uart = Uart::new();
    uart.write_str("k-os: booted on RISC-V\n");

    boot::hardware::init(&uart, hart);
    // Held until the boot hart first enters userspace; harts started
    // below wait for it. Not before init: zeroing .bss resets the lock.
    smp::lock();
    let mut blk =
        boot::hardware::probe_block_device(&uart).unwrap_or_else(|| halt(&uart, "no block device"));

    let (init, init_thread) = boot::loader::load_init(&uart, &mut blk);
    if let Some(ctx) = kern_mut().thread_context_mut(init_thread) {
        ctx.registers_mut()[10] = SHARED_BUF_VA;
    }
    let storage = boot::loader::load_server(&uart, &mut blk, b"storage.elf", "storage");
    let display = boot::loader::load_server(&uart, &mut blk, b"display.elf", "display");
    let input = boot::loader::load_server(&uart, &mut blk, b"input.elf", "input");
    let virtio_base = blk.base();

    let board = board::devices().unwrap_or_else(|| halt(&uart, "bad board table"));
    kern_mut()
        .register_board(init.pid, &board)
        .unwrap_or_else(|_| halt(&uart, "board registration failed"));
    devices::storage_setup::configure(&uart, init.pid, storage.pid, virtio_base);
    devices::gpu_setup::configure(&uart, display.pid);
    devices::input_setup::configure(&uart, input.pid);
    boot::caps::grant_all(&uart, init.pid, storage.pid, display.pid, input.pid);
    // After grant_all, so the slots init already relies on stay put.
    kern_mut()
        .grant_boot_authority(init.pid, &board)
        .unwrap_or_else(|_| halt(&uart, "boot authority failed"));

    start_other_harts(&uart, hart);
    uart.write_str("k-os: entering userspace\n");
    trap::enter_first();
}

fn start_other_harts(uart: &Uart, boot: usize) {
    unsafe extern "C" {
        fn _secondary_start();
    }
    // Harts start with the MMU off, so at the entry's physical address.
    let entry = mmu::virt_to_phys(_secondary_start as *const () as usize);
    for hart in (0..smp::num_harts()).filter(|h| *h != boot) {
        if !smp::sbi_hart_start(hart, entry) {
            halt(uart, "could not start a hart");
        }
    }
    uart.write_str("k-os: ");
    uart.write_dec(smp::num_harts());
    uart.write_str(" harts\n");
}

/// Where every hart but the boot hart starts, on its own stack.
#[unsafe(no_mangle)]
extern "C" fn kmain_secondary(_hart: usize) -> ! {
    trap::init();
    hw::timer::init();
    mmu::activate(mmu::page_table::table_at(smp::kernel_root()));
    smp::lock();
    trap::enter_first();
}

pub fn halt(uart: &Uart, msg: &str) -> ! {
    uart.write_str("k-os: FATAL — ");
    uart.write_str(msg);
    uart.put_char(b'\n');
    loop {
        unsafe { core::arch::asm!("wfi") };
    }
}

#[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") };
    }
}