use super::{latch::Latch, unwind};
use std::{any::Any, cell::UnsafeCell, mem, sync::Arc};

pub(super) enum JobResult<T> {
    None,
    Ok(T),
    Panic(Box<dyn Any + Send>),
}

/// A `Job` is used to advertise work for other threads that they may
/// want to steal. In accordance with time honored tradition, jobs are
/// arranged in a deque, so that thieves can take from the top of the
/// deque while the main worker manages the bottom of the deque. This
/// deque is managed by the `thread_pool` module.
pub(super) trait Job {
    /// Unsafe: this may be called from a different thread than the one
    /// which scheduled the job, so the implementer must ensure the
    /// appropriate traits are met, whether `Send`, `Sync`, or both.
    unsafe fn execute(this: *const ());
}

/// Effectively a Job trait object. Each JobRef **must** be executed
/// exactly once, or else data may leak.
///
/// Internally, we store the job's data in a `*const ()` pointer.  The
/// true type is something like `*const StackJob<...>`, but we hide
/// it. We also carry the "execute fn" from the `Job` trait.
pub(super) struct JobRef {
    pointer: *const (),
    execute_fn: unsafe fn(*const ()),
}

// SAFETY: every constructor of a job reference requires its job to be safe to execute on another
// thread, which is all a reference does with its pointer.
unsafe impl Send for JobRef {}
// SAFETY: a shared reference only exposes the pointer's identity, never the job.
unsafe impl Sync for JobRef {}

impl JobRef {
    /// Unsafe: caller asserts that `data` will remain valid until the
    /// job is executed.
    pub(super) unsafe fn new<T>(data: *const T) -> Self
    where
        T: Job,
    {
        // erase types:
        Self {
            pointer: data as *const (),
            execute_fn: <T as Job>::execute,
        }
    }

    /// Returns an opaque handle that can be saved and compared,
    /// without making `JobRef` itself `Copy + Eq`.
    #[inline]
    pub(super) fn id(&self) -> impl Eq + use<> {
        (self.pointer, self.execute_fn)
    }

    #[inline]
    pub(super) unsafe fn execute(self) {
        // SAFETY: the caller upholds the contract of the job this reference was built from, and
        // consuming the reference executes that job once.
        unsafe { (self.execute_fn)(self.pointer) }
    }
}

/// A job that will be owned by a stack slot. This means that when it
/// executes it need not free any heap data, the cleanup occurs when
/// the stack frame is later popped.  The function parameter indicates
/// `true` if the job was stolen and executed on a different thread.
pub(super) struct StackJob<L, F, R>
where
    L: Latch + Sync,
    F: FnOnce(bool) -> R + Send,
    R: Send,
{
    pub(super) latch: L,
    func: UnsafeCell<Option<F>>,
    result: UnsafeCell<JobResult<R>>,
}

impl<L, F, R> StackJob<L, F, R>
where
    L: Latch + Sync,
    F: FnOnce(bool) -> R + Send,
    R: Send,
{
    pub(super) const fn new(func: F, latch: L) -> Self {
        Self {
            latch,
            func: UnsafeCell::new(Some(func)),
            result: UnsafeCell::new(JobResult::None),
        }
    }

    pub(super) unsafe fn as_job_ref(&self) -> JobRef {
        // SAFETY: the caller keeps this job in place until it has executed or been run inline.
        unsafe { JobRef::new(self) }
    }

    pub(super) unsafe fn run_inline(self, stolen: bool) -> R {
        self.func.into_inner().unwrap()(stolen)
    }

    pub(super) unsafe fn into_result(self) -> R {
        self.result.into_inner().into_return_value()
    }
}

impl<L, F, R> Job for StackJob<L, F, R>
where
    L: Latch + Sync,
    F: FnOnce(bool) -> R + Send,
    R: Send,
{
    unsafe fn execute(this: *const ()) {
        // SAFETY: `this` came from `as_job_ref` on a job its owner keeps in place until the
        // latch is set, and a job reference executes once, so nothing else touches the cells.
        unsafe {
            let this = &*(this as *const Self);
            let abort = unwind::AbortIfPanic;
            let func = (*this.func.get()).take().unwrap();
            (*this.result.get()) = JobResult::call(func);
            Latch::set(&this.latch);
            mem::forget(abort);
        }
    }
}

/// Represents a job stored in the heap. Used to implement
/// `scope`. Unlike `StackJob`, when executed, `HeapJob` simply
/// invokes a closure, which then triggers the appropriate logic to
/// signal that the job executed.
///
/// (Probably `StackJob` should be refactored in a similar fashion.)
pub(super) struct HeapJob<BODY>
where
    BODY: FnOnce() + Send,
{
    job: BODY,
}

impl<BODY> HeapJob<BODY>
where
    BODY: FnOnce() + Send,
{
    pub(super) fn new(job: BODY) -> Box<Self> {
        Box::new(Self { job })
    }

    /// Creates a `JobRef` from this job. This hides all lifetimes, so
    /// it is up to you to ensure that this JobRef doesn't outlive any
    /// data that it closes over.
    pub(super) unsafe fn into_job_ref(self: Box<Self>) -> JobRef {
        // SAFETY: the box is leaked here and reclaimed by `execute`, and the caller keeps the
        // data the job closes over alive until then.
        unsafe { JobRef::new(Box::into_raw(self)) }
    }

    /// Creates a static `JobRef` from this job.
    pub(super) fn into_static_job_ref(self: Box<Self>) -> JobRef
    where
        BODY: 'static,
    {
        // SAFETY: a `'static` body closes over nothing that can expire.
        unsafe { self.into_job_ref() }
    }
}

impl<BODY> Job for HeapJob<BODY>
where
    BODY: FnOnce() + Send,
{
    unsafe fn execute(this: *const ()) {
        // SAFETY: `this` is the box leaked by `into_job_ref`, and a job reference executes once.
        let this = unsafe { Box::from_raw(this as *mut Self) };
        (this.job)();
    }
}

/// Represents a job stored in an `Arc`. Like `HeapJob`, but may be
/// turned into multiple `JobRef`s and called multiple times.
pub(super) struct ArcJob<BODY>
where
    BODY: Fn() + Send + Sync,
{
    job: BODY,
}

impl<BODY> ArcJob<BODY>
where
    BODY: Fn() + Send + Sync,
{
    pub(super) fn new(job: BODY) -> Arc<Self> {
        Arc::new(Self { job })
    }

    /// Creates a `JobRef` from this job. This hides all lifetimes, so
    /// it is up to you to ensure that this JobRef doesn't outlive any
    /// data that it closes over.
    pub(super) unsafe fn as_job_ref(this: &Arc<Self>) -> JobRef {
        // SAFETY: each reference owns one count of the `Arc`, released by `execute`, and the
        // caller keeps the data the job closes over alive until then.
        unsafe { JobRef::new(Arc::into_raw(Arc::clone(this))) }
    }

    /// Creates a static `JobRef` from this job.
    pub(super) fn as_static_job_ref(this: &Arc<Self>) -> JobRef
    where
        BODY: 'static,
    {
        // SAFETY: a `'static` body closes over nothing that can expire.
        unsafe { Self::as_job_ref(this) }
    }
}

impl<BODY> Job for ArcJob<BODY>
where
    BODY: Fn() + Send + Sync,
{
    unsafe fn execute(this: *const ()) {
        // SAFETY: `this` holds the count taken by `as_job_ref`, and a job reference executes once.
        let this = unsafe { Arc::from_raw(this as *mut Self) };
        (this.job)();
    }
}

impl<T> JobResult<T> {
    fn call(func: impl FnOnce(bool) -> T) -> Self {
        match unwind::halt_unwinding(|| func(true)) {
            Ok(x) => Self::Ok(x),
            Err(x) => Self::Panic(x),
        }
    }

    /// Convert the `JobResult` for a job that has finished (and hence
    /// its JobResult is populated) into its return value.
    ///
    /// NB. This will panic if the job panicked.
    pub(super) fn into_return_value(self) -> T {
        match self {
            Self::None => unreachable!(),
            Self::Ok(x) => x,
            Self::Panic(x) => unwind::resume_unwinding(x),
        }
    }
}