//! Threads (SYS_THREAD_SPAWN / SYS_THREAD_JOIN).
//!
//! The kernel uses the caller's stack verbatim and delivers one argument in
//! %rdi. A thread cannot free the stack it runs on, so the [`JoinHandle`]
//! owns it and frees it after a successful join; dropping an unjoined
//! handle leaks the stack (the thread may still be running on it).

use alloc::alloc::{alloc, dealloc, Layout};
use alloc::boxed::Box;

use crate::sys::{self, Status};

pub const DEFAULT_STACK_SIZE: usize = 256 * 1024;
const STACK_ALIGN: usize = 16;

pub struct JoinHandle {
    tid: u32,
    stack: *mut u8,
    stack_layout: Layout,
}

unsafe impl Send for JoinHandle {}

impl JoinHandle {
    pub fn id(&self) -> u32 {
        self.tid
    }

    /// Wait for the thread to finish and return its exit code.
    pub fn join(self) -> Result<i32, Status> {
        let mut code: i32 = 0;
        let rc = unsafe {
            sys::syscall2(sys::SYS_THREAD_JOIN, self.tid as u64, &mut code as *mut i32 as u64)
        };
        // The kernel reports a bad tid as -1 (GenericFailure).
        Status::from_result(rc)?;
        unsafe { dealloc(self.stack, self.stack_layout) };
        Ok(code)
    }
}

/// Spawn a thread running `f`; its return value becomes the exit code.
pub fn spawn<F>(f: F) -> Result<JoinHandle, Status>
where
    F: FnOnce() -> i32 + Send + 'static,
{
    spawn_with_stack(DEFAULT_STACK_SIZE, f)
}

pub fn spawn_with_stack<F>(stack_size: usize, f: F) -> Result<JoinHandle, Status>
where
    F: FnOnce() -> i32 + Send + 'static,
{
    let stack_layout = Layout::from_size_align(stack_size.max(4096), STACK_ALIGN)
        .map_err(|_| Status::InvalidArgument)?;
    let stack = unsafe { alloc(stack_layout) };
    if stack.is_null() {
        return Err(Status::OutOfMemory);
    }
    let closure = Box::into_raw(Box::new(f));
    let top = (stack as usize + stack_layout.size()) & !(STACK_ALIGN - 1);

    let rc = unsafe {
        sys::syscall3(sys::SYS_THREAD_SPAWN, trampoline::<F> as *const () as u64, closure as u64, top as u64)
    };
    match Status::from_result(rc) {
        Ok(tid) => Ok(JoinHandle { tid: tid as u32, stack, stack_layout }),
        Err(e) => unsafe {
            drop(Box::from_raw(closure));
            dealloc(stack, stack_layout);
            Err(e)
        },
    }
}

extern "C" fn trampoline<F>(arg: *mut F) -> !
where
    F: FnOnce() -> i32,
{
    let f = unsafe { Box::from_raw(arg) };
    let code = f();
    unsafe { sys::syscall1(sys::SYS_THREAD_EXIT, code as u32 as u64) };
    unreachable!()
}

/// The calling thread's id (equals the process id on the main thread).
pub fn current_id() -> u32 {
    unsafe { sys::syscall0(sys::SYS_THREAD_SELF) as u32 }
}

pub fn yield_now() {
    unsafe { sys::syscall0(sys::SYS_YIELD) };
}