//! Preemption hooks.
//!
//! Preemption is event-driven: every preemption event decrements a
//! thread-local counter (`ALLOC_COUNT`). When the counter hits zero, we
//! read RDTSC and, if the actor's timeslice has expired, call
//! `switch_to_scheduler` to yield. Resetting the counter to `ALLOC_INTERVAL`
//! amortises the RDTSC across many cheap events.
//!
//! Two event sources today:
//!   - `PreemptingAllocator` — heap allocations.
//!   - `smarm::check!()` — explicit preemption point for tight no-alloc
//!     loops, since stable Rust gives us no transparent way to preempt
//!     such loops (`__rust_probestack` is emitted inline by LLVM and not
//!     called at runtime).
//!
//! Both sources share `ALLOC_COUNT`, so the timeslice check fires at the
//! same rate regardless of whether the actor is alloc-heavy, check-heavy,
//! or mixed.
//!
//! All state is thread-local. The scheduler enables preemption on resume
//! and disables it on the return path, so the scheduler can never preempt
//! itself.
//!
//! TSC frequency is machine-dependent; `TIMESLICE_CYCLES` is a constant
//! calibrated for ~100µs on a 3 GHz CPU. A real implementation would
//! measure it at startup. For v0.1 the constant suffices.

use std::alloc::{GlobalAlloc, Layout, System};
use std::cell::Cell;
use std::sync::atomic::{AtomicBool, Ordering};

pub const DEFAULT_ALLOC_INTERVAL: u32 = 128;
pub const DEFAULT_TIMESLICE_CYCLES: u64 = 300_000; // ≈ 100µs on a 3 GHz CPU

thread_local! {
    /// While `false`, the allocator hook is a no-op.
    pub static PREEMPTION_ENABLED: Cell<bool> = const { Cell::new(false) };

    /// Countdown to next RDTSC check. Reset to `ALLOC_INTERVAL` on resume.
    static ALLOC_COUNT: Cell<u32> = const { Cell::new(DEFAULT_ALLOC_INTERVAL) };

    /// RDTSC value written by the scheduler on every actor resume.
    static TIMESLICE_START: Cell<u64> = const { Cell::new(0) };

    /// Per-thread copy of the configured alloc interval, written once at
    /// scheduler-thread startup. Kept in a thread-local so the hot path
    /// (`maybe_preempt`) pays only a TLS load, with no cache-coherency traffic.
    static CONFIGURED_ALLOC_INTERVAL: Cell<u32> = const { Cell::new(DEFAULT_ALLOC_INTERVAL) };

    /// Per-thread copy of the configured timeslice, written once at
    /// scheduler-thread startup.
    static CONFIGURED_TIMESLICE_CYCLES: Cell<u64> = const { Cell::new(DEFAULT_TIMESLICE_CYCLES) };

    /// Raw pointer to the on-CPU actor's cancellation flag, set by the
    /// scheduler on resume and reset to null when control returns to it.
    /// Null while no actor is running, so observation points are inert on the
    /// scheduler's own stack. A raw pointer (not a cloned `Arc`) keeps the
    /// resume path free of atomic ref-count traffic; see `check_cancelled` for
    /// the safety argument.
    static CURRENT_STOP: Cell<*const AtomicBool> = const { Cell::new(std::ptr::null()) };
}

// ---------------------------------------------------------------------------
// Cooperative-cancellation observation
// ---------------------------------------------------------------------------

/// Bind the on-CPU actor's stop flag. Called by the scheduler immediately
/// before `switch_to_actor`. `flag` must point into the resuming actor's
/// `Arc<AtomicBool>` heap box.
pub(crate) fn set_current_stop(flag: *const AtomicBool) {
    CURRENT_STOP.with(|c| c.set(flag));
}

/// Unbind the stop flag. Called by the scheduler on the return path from an
/// actor, so the pointer never leaks into the scheduler loop or the next actor.
pub(crate) fn clear_current_stop() {
    CURRENT_STOP.with(|c| c.set(std::ptr::null()));
}

/// Observation point for cooperative cancellation. If the on-CPU actor has
/// been flagged for stop, raise the sentinel panic so the trampoline's
/// `catch_unwind` tears the stack down (running Drop) and reports
/// `Outcome::Stopped`, distinct from a user `Panic`. A no-op (one TLS load,
/// one relaxed atomic load) on the common path.
///
/// Called from `maybe_preempt` (amortised, the `check!()`/alloc path) and from
/// the wakeup side of every blocking park (`park_current`/`yield_now`), which
/// is past the prep-to-park window — so it can never lose a wakeup.
#[inline]
pub fn check_cancelled() {
    let p = CURRENT_STOP.with(|c| c.get());
    // SAFETY: `p` is either null (no actor on-CPU — the scheduler clears it on
    // every return) or a pointer into the on-CPU actor's `Arc<AtomicBool>`
    // heap box. That box outlives the run: an actor's slot is not finalized or
    // reclaimed while the actor is on-CPU (finalize runs only after it yields
    // back), and the Arc keeps the box alive at a fixed address even if the
    // slot table Vec reallocates underneath it.
    if !p.is_null() && unsafe { (*p).load(Ordering::Relaxed) } {
        std::panic::panic_any(crate::actor::StopSentinel);
    }
}

/// Called once per scheduler thread at startup (before any actor runs).
/// Writes the runtime-configured preemption knobs into thread-locals so the
/// hot path reads them without any cross-thread coherency cost.
pub fn configure_preempt(alloc_interval: u32, timeslice_cycles: u64) {
    CONFIGURED_ALLOC_INTERVAL.with(|c| c.set(alloc_interval));
    CONFIGURED_TIMESLICE_CYCLES.with(|c| c.set(timeslice_cycles));
    // Also prime the countdown so the first resume uses the right interval.
    ALLOC_COUNT.with(|c| c.set(alloc_interval));
}

/// Arm the timeslice. Called by the scheduler on every resume.
pub fn reset_timeslice() {
    ALLOC_COUNT.with(|c| c.set(CONFIGURED_ALLOC_INTERVAL.with(|i| i.get())));
    TIMESLICE_START.with(|c| c.set(rdtsc()));
}

#[inline(always)]
pub fn rdtsc() -> u64 {
    unsafe {
        // SAFETY: x86-64 only. `lfence` serialises the instruction stream so
        // we don't measure time before prior instructions retire.
        core::arch::asm!("lfence", options(nostack, nomem, preserves_flags));
        core::arch::x86_64::_rdtsc()
    }
}

pub struct PreemptingAllocator;

unsafe impl GlobalAlloc for PreemptingAllocator {
    #[inline]
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        maybe_preempt();
        unsafe { System.alloc(layout) }
    }

    #[inline]
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        unsafe { System.dealloc(ptr, layout) }
    }

    #[inline]
    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        maybe_preempt();
        unsafe { System.alloc_zeroed(layout) }
    }

    #[inline]
    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        maybe_preempt();
        unsafe { System.realloc(ptr, layout, new_size) }
    }
}

/// Shared preemption check. Called by every preemption event source — the
/// heap allocator today, `smarm::check!()` for tight no-alloc loops.
/// Decrements `ALLOC_COUNT`; every `ALLOC_INTERVAL` calls reads the
/// timeslice clock and yields if expired.
///
/// **Invariant**: must not be called inside a "prep-to-park" region —
/// e.g. between registering as a channel's parked receiver and calling
/// `park_current()`. A preemption-driven yield in that window would
/// reach the scheduler with state=Runnable, the unparker would no-op,
/// the actor would then park, and the wakeup would be lost. Library
/// code that touches the parking primitives must keep its prep-to-park
/// regions allocation-free and check!()-free.
#[inline(always)]
pub fn maybe_preempt() {
    ALLOC_COUNT.with(|c| {
        let n = c.get();
        if n == 0 {
            c.set(CONFIGURED_ALLOC_INTERVAL.with(|i| i.get()));
            if PREEMPTION_ENABLED.with(|e| e.get()) {
                // Cooperative cancellation shares the amortised cadence with
                // the timeslice check, and shares its gate: while preemption
                // is disabled (`NoPreempt`, `with_shared`, channel critical
                // sections) the stop sentinel must NOT be raised, because an
                // allocation-triggered unwind inside a region holding a
                // `std::sync::Mutex` would poison it — one `request_stop` at
                // the wrong moment would then cascade `lock().unwrap()`
                // panics through every later user of that lock. Observation
                // is merely deferred to the next enabled allocation or the
                // wakeup side of the next park/yield, both of which are
                // lock-free points by construction.
                //
                // Observe a pending stop first: if we are being cancelled
                // there is no point yielding, we unwind instead.
                check_cancelled();
                let start = TIMESLICE_START.with(|s| s.get());
                if rdtsc().saturating_sub(start) > CONFIGURED_TIMESLICE_CYCLES.with(|t| t.get()) {
                    // SAFETY: reachable only inside an actor (the scheduler
                    // sets PREEMPTION_ENABLED on resume and clears it on
                    // return). The scheduler stack is therefore valid.
                    unsafe { crate::context::switch_to_scheduler() };
                }
            }
        } else {
            c.set(n - 1);
        }
    });
}

// ---------------------------------------------------------------------------
// Test helpers
// ---------------------------------------------------------------------------

/// Force-expire the timeslice so the next RDTSC check preempts.
pub fn expire_timeslice_for_test() {
    TIMESLICE_START.with(|c| c.set(0));
    ALLOC_COUNT.with(|c| c.set(0));
}