mod ecah;
mod kernel_services;
mod mem;

use {
    alloc::alloc::{alloc_zeroed, dealloc},
    core::{
        alloc::Layout,
        future::Future,
        pin::Pin,
        task::{Context, Poll},
    },
    hbvm::{
        mem::{softpaging::HandlePageFault, Address},
        VmRunError, VmRunOk,
    },
    log::error,
};

const STACK_SIZE: usize = 1024 * 1024;
const TIMER_QUOTIENT: usize = 1000;
type Vm = hbvm::Vm<mem::Memory, TIMER_QUOTIENT>;

pub struct ExecThread {
    vm: Vm,
    stack_bottom: *mut u8,
}

unsafe impl Send for ExecThread {}

impl ExecThread {
    pub fn set_arguments(&mut self, ptr: u64, length: u64) {
        self.vm.registers[1] = hbvm::value::Value(ptr);
        self.vm.registers[2] = hbvm::value::Value(length);
    }

    pub unsafe fn new(program: &[u8], entrypoint: Address) -> Self {
        let mut vm = Vm::new(
            mem::Memory {},
            Address::new(program.as_ptr() as u64 + entrypoint.get()),
        );

        let stack_bottom = allocate_stack();

        vm.write_reg(254, (stack_bottom as usize + STACK_SIZE - 1) as u64);

        ExecThread { vm, stack_bottom }
    }
}

impl<'p> Drop for ExecThread {
    fn drop(&mut self) {
        unsafe { dealloc(self.stack_bottom, stack_layout()) };
    }
}

impl<'p> Future for ExecThread {
    type Output = Result<(), VmRunError>;

    #[inline(always)]
    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        match self.vm.run() {
            Err(err) => {
                error!("HBVM Error\r\nRegister dump: {:?}", self.vm.registers);
                Poll::Ready(Err(err))
            }
            Ok(VmRunOk::End) => Poll::Ready(Ok(())),
            Ok(VmRunOk::Ecall) => {
                ecah::handler(&mut self.vm);
                cx.waker().wake_by_ref();
                Poll::Pending
            }
            Ok(VmRunOk::Timer) => {
                cx.waker().wake_by_ref();
                Poll::Pending
            }
            Ok(VmRunOk::Breakpoint) => {
                error!(
                    "HBVM Debug breakpoint\r\nRegister dump: {:?}",
                    self.vm.registers
                );
                cx.waker().wake_by_ref();
                Poll::Pending
            }
        }
    }
}

struct PageFaultHandler;
impl HandlePageFault for PageFaultHandler {
    fn page_fault(
        &mut self,
        reason: hbvm::mem::MemoryAccessReason,
        _pagetable: &mut hbvm::mem::softpaging::paging::PageTable,
        vaddr: hbvm::mem::Address,
        size: hbvm::mem::softpaging::PageSize,
        dataptr: *mut u8,
    ) -> bool {
        error!("REASON: {reason} vaddr: {vaddr} size: {size:?} Dataptr {dataptr:p}");
        false
    }
}

#[inline(always)]
const fn stack_layout() -> Layout {
    unsafe { Layout::from_size_align_unchecked(STACK_SIZE, 4096) }
}

#[inline(always)]
fn allocate_stack() -> *mut u8 {
    unsafe { alloc_zeroed(stack_layout()) }
}