#![no_std]
#![no_main]

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

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

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();
    let pool = domain::memory::page_pool::PagePool::new(
        types::address::PhysicalAddress::from_raw(board::USER_POOL_BASE),
        board::USER_POOL_PAGES,
    );
    let mut kern = kernel::Kernel::new(riscv_platform, pool);
    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 (pid=");
    uart.write_dec(pid.raw() as usize);
    uart.write_str(", tid=");
    uart.write_dec(tid.raw() as usize);
    uart.write_str(")\n");

    ipc_demo(&mut kern, &uart, pid);

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

fn ipc_demo(
    kern: &mut kernel::Kernel<platform::RiscVPlatform>,
    uart: &Uart,
    init_pid: types::id::ProcessId,
) {
    use domain::message::Message;
    use domain::resource::Resource;
    use domain::syscall::Syscall;
    use kernel::{SyscallOutcome, SyscallResult};
    use types::address::DataRef;
    use types::bounded::CapSlot;
    use types::id::ThreadId;
    use types::permission::{Permission, PermissionSet};

    uart.write_str("k-os: --- IPC demo ---\n");

    let _proc_b_slot = match kern.handle_syscall(Syscall::CreateProcess {
        create_cap: CapSlot::new(0).unwrap_or_else(|| unreachable!()),
        image: DataRef::empty(),
    }) {
        SyscallOutcome::Complete(SyscallResult::Slot(s)) => s,
        _ => {
            uart.write_str("k-os: FAIL creating process B\n");
            return;
        }
    };
    let pid_b = types::id::ProcessId::from_raw(1);
    let tid_b = ThreadId::from_raw(1);
    uart.write_str("k-os: process B created\n");

    let ep_id = kern.create_endpoint().unwrap_or_else(|_| unreachable!());

    let ep_cap_a = kern.grant_capability(
        init_pid,
        Resource::Endpoint(ep_id),
        PermissionSet::empty()
            .with(Permission::Execute)
            .with(Permission::Read),
    ).unwrap_or_else(|_| unreachable!());

    let ep_cap_b = kern.grant_capability(
        pid_b,
        Resource::Endpoint(ep_id),
        PermissionSet::empty()
            .with(Permission::Execute)
            .with(Permission::Read),
    ).unwrap_or_else(|_| unreachable!());

    uart.write_str("k-os: endpoint created, caps granted\n");

    kern.switch_to_thread(tid_b);
    match kern.handle_syscall(Syscall::Receive { endpoint: ep_cap_b }) {
        SyscallOutcome::BlockCaller(_) => {
            uart.write_str("k-os: process B blocked on Receive\n");
        }
        _ => {
            uart.write_str("k-os: UNEXPECTED outcome for Receive\n");
            return;
        }
    }

    let tid_a = ThreadId::from_raw(0);
    kern.switch_to_thread(tid_a);
    let msg = Message::from_words([42, 7, 0, 0]);
    match kern.handle_syscall(Syscall::Send {
        endpoint: ep_cap_a,
        message: msg,
    }) {
        SyscallOutcome::SwitchTo(target) => {
            uart.write_str("k-os: Send delivered, switching to thread ");
            uart.write_dec(target.raw() as usize);
            uart.put_char(b'\n');
        }
        _ => {
            uart.write_str("k-os: UNEXPECTED outcome for Send\n");
            return;
        }
    }

    match kern.pending_result(tid_b) {
        Some(SyscallResult::Message(m)) => {
            uart.write_str("k-os: IPC complete — received: [");
            uart.write_dec(m.word(0));
            uart.write_str(", ");
            uart.write_dec(m.word(1));
            uart.write_str(", ");
            uart.write_dec(m.word(2));
            uart.write_str(", ");
            uart.write_dec(m.word(3));
            uart.write_str("]\n");
        }
        _ => {
            uart.write_str("k-os: FAIL — no message in pending result\n");
        }
    }
}

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