mod ecah;
mod mem;

use {
    crate::{arch, ipc::buffer::IpcBuffer, kmain::IPC_BUFFERS},
    alloc::boxed::Box,
    core::{default, future::Future, marker::PhantomData, ptr::NonNull, task::Poll},
    hbvm::{
        mem::{
            softpaging::{icache::ICache, HandlePageFault, SoftPagedMem},
            Address, Memory,
        },
        VmRunError, VmRunOk,
    },
    log::{debug, error, info, trace, warn},
};

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

pub struct ExecThread<'p> {
    vm: Vm,
    stack_top: *mut u8,
    _phantom: PhantomData<&'p [u8]>,
}

unsafe impl<'p> Send for ExecThread<'p> {}
impl<'p> ExecThread<'p> {
    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: &'p [u8], entrypoint: Address) -> Self {
        let mut vm = unsafe {
            Vm::new(
                mem::Memory,
                Address::new(program.as_ptr() as u64 + entrypoint.get()),
            )
        };

        let stack_top = unsafe { allocate_stack().as_ptr().add(STACK_SIZE - 1) };
        vm.write_reg(254, stack_top as u64);

        ExecThread {
            vm,
            stack_top,
            _phantom: Default::default(),
        }
    }
}

impl<'p> Drop for ExecThread<'p> {
    fn drop(&mut self) {
        unsafe { alloc::alloc::dealloc(self.stack_top.sub(STACK_SIZE - 1), stack_layout()) };
    }
}

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

    fn poll(
        mut self: core::pin::Pin<&mut Self>,
        cx: &mut core::task::Context<'_>,
    ) -> Poll<Self::Output> {
        match self.vm.run() {
            Err(err) => {
                log::error!("HBVM Error\r\nRegister dump: {:?}", self.vm.registers,);
                return Poll::Ready(Err(err));
            }
            Ok(VmRunOk::End) => return Poll::Ready(Ok(())),
            Ok(VmRunOk::Ecall) => ecah::handler(&mut self.vm),
            Ok(VmRunOk::Timer) => (),
            Ok(VmRunOk::Breakpoint) => {
                log::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
    where
        Self: Sized,
    {
        log::error!(
            "REASON: {reason} \
             vaddr: {vaddr} \
             size: {size:?} \
             Dataptr {dataptr:p}",
        );

        false
    }
}

const fn stack_layout() -> core::alloc::Layout {
    unsafe { alloc::alloc::Layout::from_size_align_unchecked(STACK_SIZE, 4096) }
}

fn allocate_stack() -> NonNull<u8> {
    let layout = stack_layout();
    match NonNull::new(unsafe { alloc::alloc::alloc_zeroed(layout) }) {
        Some(ptr) => ptr,
        None => alloc::alloc::handle_alloc_error(layout),
    }
}