// worker.rs --- T M3.1 work-stealing pool with cooperative cancellation.

//! Worker pool primitive (T M3.1).
//!
//! Spec contract: a work-stealing thread pool built on
//! [`std::thread`] and [`crossbeam_deque`] --- explicitly **not** a
//! generic async runtime ([spec §3 checkpoint 5]). Each worker owns
//! a local FIFO deque and steals from siblings + a shared injector
//! when its own deque is empty. Idle workers park on a [`Condvar`];
//! [`WorkerPool::dispatch`] notifies one. Dropping the pool signals
//! every worker and joins.
//!
//! # Cancellation
//!
//! Each [`JobHandle`] carries a [`CancellationToken`] (an
//! [`Arc<AtomicBool>`]). A user closure that loops on long work
//! polls the token at granular boundaries; calling
//! [`JobHandle::cancel`] flips the bit so the next check returns.
//! Cancellation is cooperative: the runtime never preempts a job.
//! A job that was cancelled before any worker picked it up sees the
//! flag at the start of execution and returns immediately without
//! invoking the user closure.
//!
//! # Panic isolation
//!
//! Each job runs inside [`std::panic::catch_unwind`]. A panicking
//! job does not kill its worker; subsequent jobs run on the same
//! thread.
//!
//! # What this layer does *not* do
//!
//! No message bus, no result delivery channel, no Lua surface ---
//! T M3.1 is the raw primitive. Result delivery is a per-call
//! responsibility: the user closure can capture a
//! [`crossbeam_channel::Sender`] and send. The typed message bus
//! (T M3.2) and the coroutine-based async API (T M3.3) build on
//! top.

use std::panic::{AssertUnwindSafe, catch_unwind};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Condvar, Mutex};
use std::thread::{self, JoinHandle};
use std::time::Duration;

use crossbeam::deque::{Injector, Steal, Stealer, Worker};

/// Stable identifier for a dispatched job, monotonically increasing
/// per [`WorkerPool`] from `0`.
pub type JobId = u64;

// ---------------------------------------------------------------------------
// CancellationToken
// ---------------------------------------------------------------------------

/// Cooperative cancellation flag shared between a [`JobHandle`] and
/// the worker closure that owns it.
///
/// `is_cancelled` is the user-side check; `cancel` is the producer
/// side. Both are lock-free atomic reads/writes. Cloning a token
/// shares the same flag --- two clones see each other's state.
#[derive(Clone, Debug, Default)]
pub struct CancellationToken {
    flag: Arc<AtomicBool>,
}

impl CancellationToken {
    /// A fresh, not-yet-cancelled token.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Has [`Self::cancel`] been called on this token (or any clone
    /// of it)? Returns `true` from then on.
    #[must_use]
    pub fn is_cancelled(&self) -> bool {
        self.flag.load(Ordering::Relaxed)
    }

    /// Mark the token cancelled. Idempotent.
    pub fn cancel(&self) {
        self.flag.store(true, Ordering::Relaxed);
    }
}

// ---------------------------------------------------------------------------
// JobHandle
// ---------------------------------------------------------------------------

/// Handle returned by [`WorkerPool::dispatch`].
///
/// Carries the job's id and the cancellation token. Dropping a
/// handle does **not** cancel the job (the user may have other
/// clones of the token); call [`Self::cancel`] explicitly.
#[derive(Clone, Debug)]
pub struct JobHandle {
    id: JobId,
    token: CancellationToken,
}

impl JobHandle {
    /// The dispatch-order id assigned by the pool.
    #[must_use]
    pub fn id(&self) -> JobId {
        self.id
    }

    /// Request cooperative cancellation. The worker closure (if it
    /// polls the token) sees the flag on its next check.
    pub fn cancel(&self) {
        self.token.cancel();
    }

    /// Has the job been cancelled?
    #[must_use]
    pub fn is_cancelled(&self) -> bool {
        self.token.is_cancelled()
    }

    /// Borrow a clone of the token, e.g. to register cancellation
    /// behaviour outside the original closure.
    #[must_use]
    pub fn token(&self) -> CancellationToken {
        self.token.clone()
    }
}

// ---------------------------------------------------------------------------
// WorkerPool
// ---------------------------------------------------------------------------

type Job = Box<dyn FnOnce() + Send + 'static>;

struct PoolShared {
    injector: Injector<Job>,
    stealers: Vec<Stealer<Job>>,
    shutdown: AtomicBool,
    next_id: AtomicU64,
    /// Idle-worker park lot. Workers acquire the mutex, recheck for
    /// work, and `wait_timeout` on the condvar. Producers and
    /// shutdown signal `notify_one` / `notify_all`. The 100ms
    /// timeout is a belt-and-braces guard against a missed
    /// notification --- workers re-poll the queues even without an
    /// explicit wakeup.
    parker: (Mutex<()>, Condvar),
}

impl PoolShared {
    fn alloc_id(&self) -> JobId {
        self.next_id.fetch_add(1, Ordering::Relaxed)
    }

    fn notify_one(&self) {
        // Holding the lock during `notify_one` is required to avoid
        // a lost wakeup when a worker is between the empty-queue
        // check and the `wait_timeout` call.
        let _guard = self.parker.0.lock().expect("parker mutex");
        self.parker.1.notify_one();
    }

    fn notify_all(&self) {
        let _guard = self.parker.0.lock().expect("parker mutex");
        self.parker.1.notify_all();
    }
}

/// Work-stealing thread pool with cooperative cancellation tokens.
///
/// Construction spawns `size` OS threads, each running its own
/// [`worker_loop`]. Drop the pool to stop them: [`Drop`] signals
/// every worker and joins.
///
/// `WorkerPool` is `Send + Sync`: dispatch can come from any
/// thread, but typical pmacs use is single-producer (the Lua main
/// thread).
pub struct WorkerPool {
    shared: Arc<PoolShared>,
    /// Join handles, drained exactly once by [`Self::shutdown`]
    /// (directly or via `Drop`). Behind a `Mutex` so shutdown works
    /// from a shared reference: the pool's owner is typically an
    /// `Rc<AsyncRuntime>` cloned into Lua closures, and those clones
    /// form VM reference cycles that keep the `Rc` from ever
    /// reaching zero --- an embedder that merely *drops* its handle
    /// would leak every worker thread. `EditorState::drop` calls
    /// `shutdown()` explicitly instead.
    workers: Mutex<Vec<JoinHandle<()>>>,
    /// Thread count at construction (stable across shutdown).
    size: usize,
}

impl WorkerPool {
    /// Build a pool with exactly `size` worker threads. `size` is
    /// clamped to at least `1` so callers that compute `cores - 1`
    /// on a 1-core machine don't end up with zero workers.
    #[must_use]
    pub fn new(size: usize) -> Self {
        let size = size.max(1);
        let local_queues: Vec<Worker<Job>> = (0..size).map(|_| Worker::new_fifo()).collect();
        let stealers: Vec<Stealer<Job>> = local_queues.iter().map(Worker::stealer).collect();
        let shared = Arc::new(PoolShared {
            injector: Injector::new(),
            stealers,
            shutdown: AtomicBool::new(false),
            next_id: AtomicU64::new(0),
            parker: (Mutex::new(()), Condvar::new()),
        });
        let workers: Vec<JoinHandle<()>> = local_queues
            .into_iter()
            .enumerate()
            .map(|(idx, local)| {
                let shared = Arc::clone(&shared);
                thread::Builder::new()
                    .name(format!("pmacs-worker-{idx}"))
                    .spawn(move || worker_loop(&local, &shared))
                    .expect("spawn worker thread")
            })
            .collect();
        Self {
            shared,
            size,
            workers: Mutex::new(workers),
        }
    }

    /// Build a pool sized at `available_parallelism - 1`, with a
    /// floor of `1`. Reserves one core for the main editor thread,
    /// matching the spec's "main thread is the event loop" pattern
    /// ([spec §6.1]).
    #[must_use]
    pub fn with_default_size() -> Self {
        let cores = thread::available_parallelism().map_or(2, std::num::NonZeroUsize::get);
        Self::new(cores.saturating_sub(1))
    }

    /// Number of worker threads this pool was built with.
    #[must_use]
    pub fn size(&self) -> usize {
        self.size
    }

    /// Signal every worker to exit, without joining. Idle (parked)
    /// workers observe the flag within their 100ms park timeout and
    /// return; a worker mid-job exits when its job finishes. Queued
    /// jobs that haven't been picked up are dropped without running;
    /// jobs dispatched *after* the signal are never picked up.
    ///
    /// Exists as an explicit method (not just `Drop`) because the
    /// pool's owning `Rc<AsyncRuntime>` is captured into Lua-VM
    /// reference cycles and may never be reclaimed --- callers that
    /// know the editor is going away (`EditorState::drop`) signal the
    /// threads down regardless.
    ///
    /// Deliberately does NOT join: a worker can be blocked publishing
    /// its reply onto the message bus that only the *main thread*
    /// drains, so a main-thread join here is a deadlock (observed as
    /// the m4 acceptance suite wedging for hours at teardown). Callers
    /// that own the whole world and want the join use
    /// [`Self::shutdown`] (or just drop the pool).
    pub fn signal_shutdown(&self) {
        self.shared.shutdown.store(true, Ordering::Release);
        self.shared.notify_all();
    }

    /// [`Self::signal_shutdown`] plus a join of every worker thread.
    /// Idempotent: the second call finds no handles and returns
    /// immediately. Only safe where no worker can be blocked on the
    /// caller's own thread (see `signal_shutdown`); `Drop` uses it
    /// because a pool being dropped has no live bus consumer to
    /// deadlock against in the bare-pool case.
    pub fn shutdown(&self) {
        self.signal_shutdown();
        let handles: Vec<JoinHandle<()>> =
            std::mem::take(&mut *self.workers.lock().expect("worker pool mutex poisoned"));
        for handle in handles {
            let _ = handle.join();
        }
    }

    /// Submit `work` to be run on a worker. Returns a [`JobHandle`]
    /// that owns the cancellation token for this job.
    ///
    /// `work` is invoked with a borrow of the same token the handle
    /// holds; the closure is responsible for polling
    /// [`CancellationToken::is_cancelled`] at granular boundaries
    /// inside any long-running loop. A job whose token is set
    /// before the worker picks it up sees the flag immediately and
    /// returns without running the user closure.
    ///
    /// Result delivery is up to the closure: capture a
    /// `crossbeam_channel::Sender` (or any other `Send` channel) to
    /// pass values back. T M3.2 will provide the message bus that
    /// formalises this; T M3.1 deliberately does not.
    pub fn dispatch<F>(&self, work: F) -> JobHandle
    where
        F: FnOnce(&CancellationToken) + Send + 'static,
    {
        let id = self.shared.alloc_id();
        let token = CancellationToken::new();
        let token_for_job = token.clone();
        let job: Job = Box::new(move || {
            // Skip user work entirely if cancelled before we ran.
            if token_for_job.is_cancelled() {
                return;
            }
            // Panic isolation: a panicking job must not poison the
            // worker thread. We swallow the panic payload here ---
            // higher layers (M3.2 message bus) will deliver
            // structured failure to the caller.
            let _ = catch_unwind(AssertUnwindSafe(|| work(&token_for_job)));
        });
        self.shared.injector.push(job);
        self.shared.notify_one();
        JobHandle { id, token }
    }
}

impl Drop for WorkerPool {
    /// Dropping the pool is an implicit [`Self::shutdown`]: every
    /// worker is signalled to exit at its next idle wakeup and
    /// joined. No-op when `shutdown` already ran.
    fn drop(&mut self) {
        self.shutdown();
    }
}

fn worker_loop(local: &Worker<Job>, shared: &Arc<PoolShared>) {
    loop {
        if shared.shutdown.load(Ordering::Acquire) {
            return;
        }
        if let Some(job) = find_work(local, shared) {
            job();
            continue;
        }
        // No work --- park briefly, then retry. The 100ms timeout
        // bounds how long a missed wakeup can stall this worker.
        let guard = shared.parker.0.lock().expect("parker mutex");
        if shared.shutdown.load(Ordering::Acquire) {
            return;
        }
        // The injector was empty when we last looked; if a producer
        // pushed in the meantime they will have signalled the
        // condvar. We hold the parker mutex now, so any signal that
        // happens after we drop the lock will wake us via wait.
        let _ = shared
            .parker
            .1
            .wait_timeout(guard, Duration::from_millis(100))
            .expect("parker condvar");
    }
}

fn find_work(local: &Worker<Job>, shared: &PoolShared) -> Option<Job> {
    if let Some(job) = local.pop() {
        return Some(job);
    }
    // Pull a batch from the global injector into our local deque,
    // returning one to run immediately.
    loop {
        match shared.injector.steal_batch_and_pop(local) {
            Steal::Success(job) => return Some(job),
            Steal::Empty => break,
            Steal::Retry => {}
        }
    }
    // Steal one job from each sibling.
    for stealer in &shared.stealers {
        loop {
            match stealer.steal() {
                Steal::Success(job) => return Some(job),
                Steal::Empty => break,
                Steal::Retry => {}
            }
        }
    }
    None
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use crossbeam::channel;
    use std::sync::atomic::AtomicU64;

    fn assert_recv_within<T>(rx: &channel::Receiver<T>, label: &str) -> T {
        rx.recv_timeout(Duration::from_secs(2))
            .unwrap_or_else(|_| panic!("timed out waiting for {label}"))
    }

    /// Acceptance bullet: jobs run on workers; results returned via
    /// callback (the user closure forwards through a channel).
    #[test]
    fn dispatch_runs_user_closure_on_a_worker() {
        let pool = WorkerPool::new(2);
        let (tx, rx) = channel::bounded::<u64>(1);
        let _ = pool.dispatch(move |_| {
            tx.send(42).unwrap();
        });
        assert_eq!(assert_recv_within(&rx, "dispatched value"), 42);
    }

    /// Job ids are unique and monotonic per pool.
    #[test]
    fn job_ids_are_monotonic() {
        let pool = WorkerPool::new(1);
        let h1 = pool.dispatch(|_| {});
        let h2 = pool.dispatch(|_| {});
        let h3 = pool.dispatch(|_| {});
        assert!(h1.id() < h2.id() && h2.id() < h3.id());
    }

    /// A job whose token is set before the worker picks it up must
    /// not run the user closure. We park the only worker on a
    /// blocking job, then dispatch the doomed job, cancel it,
    /// release the worker, and verify the doomed closure never set
    /// its sentinel flag.
    #[test]
    fn cancel_before_dispatch_skips_user_work() {
        let pool = WorkerPool::new(1);
        let (release, gate) = channel::bounded::<()>(0);
        let _h_block = pool.dispatch(move |_| {
            let _ = gate.recv();
        });
        let did_run = Arc::new(AtomicBool::new(false));
        let did_run_clone = Arc::clone(&did_run);
        let h_doomed = pool.dispatch(move |_| {
            did_run_clone.store(true, Ordering::SeqCst);
        });
        h_doomed.cancel();
        release.send(()).unwrap();
        // Give the worker time to drain and run the doomed job.
        thread::sleep(Duration::from_millis(50));
        assert!(
            !did_run.load(Ordering::SeqCst),
            "cancelled job should not have run"
        );
    }

    /// A long-running job that polls the token sees cancellation
    /// observed mid-flight.
    #[test]
    fn cancel_during_work_observed_by_user_closure() {
        let pool = WorkerPool::new(1);
        let started = Arc::new(AtomicBool::new(false));
        let observed = Arc::new(AtomicBool::new(false));
        let started_c = Arc::clone(&started);
        let observed_c = Arc::clone(&observed);
        let h = pool.dispatch(move |t| {
            started_c.store(true, Ordering::SeqCst);
            while !t.is_cancelled() {
                thread::sleep(Duration::from_millis(1));
            }
            observed_c.store(true, Ordering::SeqCst);
        });
        // Wait for the worker to enter the loop.
        let deadline = std::time::Instant::now() + Duration::from_secs(2);
        while !started.load(Ordering::SeqCst) {
            assert!(
                std::time::Instant::now() < deadline,
                "worker never started job"
            );
            thread::sleep(Duration::from_millis(1));
        }
        h.cancel();
        let deadline = std::time::Instant::now() + Duration::from_secs(2);
        while !observed.load(Ordering::SeqCst) {
            assert!(
                std::time::Instant::now() < deadline,
                "worker never observed cancellation"
            );
            thread::sleep(Duration::from_millis(1));
        }
    }

    /// A panicking job must not kill its worker. After a panic on a
    /// single-thread pool, a follow-up job must still run on the
    /// (same) worker.
    #[test]
    fn panic_in_job_does_not_kill_worker() {
        let pool = WorkerPool::new(1);
        let _ = pool.dispatch(|_| panic!("boom"));
        // Allow the worker to swallow the panic.
        thread::sleep(Duration::from_millis(20));
        let (tx, rx) = channel::bounded::<u32>(1);
        let _ = pool.dispatch(move |_| {
            tx.send(7).unwrap();
        });
        assert_eq!(assert_recv_within(&rx, "post-panic dispatch"), 7);
    }

    /// `with_default_size` produces at least one worker on every
    /// platform, even single-core machines.
    #[test]
    fn default_size_floors_at_one() {
        let pool = WorkerPool::with_default_size();
        assert!(
            pool.size() >= 1,
            "default size must be at least 1, got {}",
            pool.size()
        );
    }

    /// Acceptance bullet: 10000 dispatches with random
    /// cancellations, no leaks or hangs. We use a deterministic
    /// "cancel every Nth" pattern instead of an RNG so the test is
    /// reproducible. Drop is allowed to discard queued work after
    /// setting shutdown; this test verifies the stress path does not
    /// hang, panic, or over-run completions.
    #[test]
    fn stress_10k_dispatches_with_periodic_cancels_no_hang() {
        const TOTAL: usize = 10_000;
        const CANCEL_EVERY: usize = 3;
        let pool = WorkerPool::with_default_size();
        let completed = Arc::new(AtomicU64::new(0));
        let mut handles = Vec::with_capacity(TOTAL);
        for i in 0..TOTAL {
            let completed = Arc::clone(&completed);
            let h = pool.dispatch(move |t| {
                if t.is_cancelled() {
                    return;
                }
                // A bit of trivial work that observes the token.
                let mut acc: u64 = 0;
                for j in 0..32u64 {
                    if t.is_cancelled() {
                        return;
                    }
                    acc = acc.wrapping_add(j);
                }
                completed.fetch_add(1, Ordering::Relaxed);
                std::hint::black_box(acc);
            });
            if i % CANCEL_EVERY == 0 {
                h.cancel();
            }
            handles.push(h);
        }
        // Drop may discard still-queued work after setting shutdown,
        // so on a loaded CI runner it can win the race before any
        // worker finishes even one job, flaking the `count > 0`
        // assert below (observed on the macOS runner). ~2/3 of the
        // queue is non-cancelled, so one completion must land unless
        // workers are wedged — wait for it, bounded.
        let wait_start = std::time::Instant::now();
        while completed.load(Ordering::Relaxed) == 0
            && wait_start.elapsed() < std::time::Duration::from_secs(10)
        {
            std::thread::yield_now();
        }
        // Drain shutdown synchronously via Drop: this returns only
        // after every queued, non-cancelled job has run (or every
        // cancelled job has either run-then-noop or been silently
        // dropped). No hang means we get here within the test
        // harness's default timeout.
        drop(pool);
        // Cancelled-before-pickup jobs return without bumping the
        // counter; cancelled-after-pickup jobs may bump 0 or 1
        // times depending on when they polled. We only assert no
        // job over-counted: completion count must not exceed the
        // count of jobs that were ever eligible.
        let count = completed.load(Ordering::Relaxed);
        let max_eligible = u64::try_from(TOTAL).unwrap();
        assert!(
            count <= max_eligible,
            "completion count {count} exceeded eligible {max_eligible}"
        );
        assert!(
            count > 0,
            "stress test should complete at least one job before shutdown"
        );
    }

    /// Drop semantics: dropping a `WorkerPool` must join every
    /// worker. We rely on the test harness's process-exit timeout
    /// to flag failure if a worker leaks.
    #[test]
    fn drop_joins_workers() {
        let pool = WorkerPool::new(4);
        let (tx, rx) = channel::bounded::<()>(1);
        let _ = pool.dispatch(move |_| {
            tx.send(()).unwrap();
        });
        assert_recv_within(&rx, "pre-drop dispatch");
        drop(pool);
        // If Drop didn't join, the process would have to wait for
        // detached threads to exit naturally. Reaching here is the
        // assertion.
    }
}