//! The single-threaded scheduler.
//!
//! There is one global scheduler per OS thread, stored in a thread-local.
//! `run(initial)` initialises it, spawns the initial actor, drives the loop
//! until the run queue is empty, then tears it down.
//!
//! Slot table: a `Vec<Slot>` indexed by `Pid::index()`, with a free list of
//! reusable indices. Each slot has a `generation` counter that increments
//! every time the slot is freed; `Pid` carries the generation it was minted
//! with, so a stale PID has a mismatching generation and is detected on
//! lookup.
//!
//! Run queue: a `VecDeque<Pid>` of runnable actors. The state of an actor
//! is implicit in slot.state: `Runnable` means it's either in the queue or
//! currently executing; `Parked` means it's waiting for something to unpark
//! it (channel send, join completion, …); `Done` means it has finished and
//! is awaiting reaping.
//!
//! Joining: `JoinHandle::join()` parks the calling actor and registers it
//! on the target slot's `waiters` list. When the target actor finishes,
//! the scheduler reaps the slot and unparks every waiter, passing them the
//! outcome via a side channel (the target's `outcome` field, drained on
//! the joiner side).

use crate::actor::{
    clear_current_pid, current_pid, is_actor_done, reset_actor_done,
    set_current_actor_box, set_current_pid, take_last_outcome, trampoline, Actor, Outcome,
};
use crate::channel::Sender;
use crate::context::{get_actor_sp, init_actor_stack, set_actor_sp, switch_to_actor};
use crate::pid::Pid;
use crate::preempt::PREEMPTION_ENABLED;
use crate::stack::Stack;
use crate::supervisor::Signal;
use std::cell::RefCell;
use std::collections::VecDeque;

// ---------------------------------------------------------------------------
// Configuration
// ---------------------------------------------------------------------------

const ACTOR_STACK_SIZE: usize = 64 * 1024;

// ---------------------------------------------------------------------------
// Per-actor slot
// ---------------------------------------------------------------------------

enum State {
    /// Either in the run queue or currently executing.
    Runnable,
    /// Removed from the queue, waiting for `unpark()`.
    Parked,
    /// The actor has finished. Slot persists until the last `JoinHandle`
    /// has been joined (or dropped). Then the slot is freed.
    Done,
}

struct Slot {
    /// Bumped every time this slot is freed and re-used. A `Pid` with a
    /// non-matching generation is stale.
    generation: u32,
    /// `None` when the slot is free. `Some` otherwise.
    actor: Option<Actor>,
    state: State,
    /// PIDs waiting in `JoinHandle::join`.
    waiters: Vec<Pid>,
    /// The outcome the actor produced, captured when it finished.
    /// Drained by `JoinHandle::join`.
    outcome: Option<Outcome>,
    /// If this slot is a supervisor, the sender into its `Signal` mailbox.
    /// Cloned out and used when one of its children dies.
    supervisor_channel: Option<Sender<Signal>>,
    /// Number of `JoinHandle`s still outstanding for this actor. The slot
    /// is reclaimed only when the actor is done AND outstanding_handles == 0.
    outstanding_handles: u32,
}

impl Slot {
    fn vacant() -> Self {
        Self {
            generation: 0,
            actor: None,
            state: State::Done,
            waiters: Vec::new(),
            outcome: None,
            supervisor_channel: None,
            outstanding_handles: 0,
        }
    }
}

// ---------------------------------------------------------------------------
// Scheduler state
// ---------------------------------------------------------------------------

struct SchedulerState {
    slots: Vec<Slot>,
    free_list: Vec<u32>,
    run_queue: VecDeque<Pid>,
    /// The root supervisor's PID. Children spawned at the top level are
    /// supervised by this. Set by `run()`.
    root_pid: Option<Pid>,
}

impl SchedulerState {
    fn new() -> Self {
        Self {
            slots: Vec::new(),
            free_list: Vec::new(),
            run_queue: VecDeque::new(),
            root_pid: None,
        }
    }

    /// Allocate a slot; return its (index, generation).
    fn allocate_slot(&mut self) -> (u32, u32) {
        if let Some(idx) = self.free_list.pop() {
            let s = &mut self.slots[idx as usize];
            (idx, s.generation)
        } else {
            let idx = self.slots.len() as u32;
            self.slots.push(Slot::vacant());
            (idx, 0)
        }
    }

    fn slot(&self, pid: Pid) -> Option<&Slot> {
        let s = self.slots.get(pid.index() as usize)?;
        if s.generation == pid.generation() { Some(s) } else { None }
    }

    fn slot_mut(&mut self, pid: Pid) -> Option<&mut Slot> {
        let s = self.slots.get_mut(pid.index() as usize)?;
        if s.generation == pid.generation() { Some(s) } else { None }
    }
}

thread_local! {
    static SCHED: RefCell<Option<SchedulerState>> = const { RefCell::new(None) };
}

fn with_sched<R>(f: impl FnOnce(&mut SchedulerState) -> R) -> R {
    SCHED.with(|c| {
        let mut g = c.borrow_mut();
        let s = g.as_mut().expect("scheduler not running");
        f(s)
    })
}

/// Same as `with_sched` but returns `None` when there's no scheduler instead
/// of panicking. Used on cleanup paths (channel sender drop during shutdown,
/// for example).
fn try_with_sched<R>(f: impl FnOnce(&mut SchedulerState) -> R) -> Option<R> {
    SCHED.with(|c| {
        let mut g = c.borrow_mut();
        g.as_mut().map(f)
    })
}

// ---------------------------------------------------------------------------
// JoinHandle
// ---------------------------------------------------------------------------

#[derive(Debug)]
pub struct JoinError {
    /// Whatever `panic!` was called with.
    pub payload: Box<dyn std::any::Any + Send>,
}

pub struct JoinHandle {
    pid: Pid,
    /// `false` once `join()` has been called and the handle has consumed
    /// its outcome. Prevents the Drop impl from double-decrementing.
    consumed: bool,
}

impl JoinHandle {
    pub fn pid(&self) -> Pid { self.pid }

    /// Block the calling actor until the target completes. Returns
    /// `Ok(())` on normal exit, `Err(JoinError)` if the target panicked.
    pub fn join(mut self) -> Result<(), JoinError> {
        let me = current_pid().expect("join() called outside an actor");

        loop {
            let outcome = with_sched(|s| {
                let slot = s.slot_mut(self.pid)
                    .expect("join: target slot has been reused");
                if matches!(slot.state, State::Done) {
                    Some(slot.outcome.take().expect("Done slot must have an outcome"))
                } else {
                    slot.waiters.push(me);
                    None
                }
            });

            match outcome {
                Some(o) => {
                    self.consumed = true;
                    self.decrement_handle_count();
                    return match o {
                        Outcome::Exit => Ok(()),
                        Outcome::Panic(p) => Err(JoinError { payload: p }),
                    };
                }
                None => park_current(),
            }
        }
    }

    fn decrement_handle_count(&mut self) {
        with_sched(|s| {
            let should_reclaim = match s.slot_mut(self.pid) {
                Some(slot) => {
                    slot.outstanding_handles = slot.outstanding_handles.saturating_sub(1);
                    matches!(slot.state, State::Done) && slot.outstanding_handles == 0
                }
                None => false,
            };
            if should_reclaim {
                reclaim_slot(s, self.pid);
            }
        });
    }
}

impl Drop for JoinHandle {
    fn drop(&mut self) {
        if !self.consumed {
            self.decrement_handle_count();
        }
    }
}

// ---------------------------------------------------------------------------
// Slot reclamation
// ---------------------------------------------------------------------------

fn reclaim_slot(s: &mut SchedulerState, pid: Pid) {
    let idx = pid.index();
    let slot = &mut s.slots[idx as usize];
    // Bump generation so any stale PIDs from now on miss.
    slot.generation = slot.generation.wrapping_add(1);
    // Drop the actor (its stack with it).
    slot.actor = None;
    slot.outcome = None;
    slot.waiters.clear();
    slot.supervisor_channel = None;
    slot.state = State::Done; // semantically vacant; allocator checks free_list
    slot.outstanding_handles = 0;
    s.free_list.push(idx);
}

// ---------------------------------------------------------------------------
// spawn / spawn_under / self_pid
// ---------------------------------------------------------------------------

/// Spawn `f` as a child of the currently-executing actor.
/// Outside an actor (only legal from `run()`'s initial setup), the child's
/// supervisor is the root supervisor.
pub fn spawn(f: impl FnOnce() + Send + 'static) -> JoinHandle {
    let parent = current_pid()
        .or_else(|| with_sched(|s| s.root_pid))
        .expect("spawn() before run()");
    spawn_under(parent, f)
}

/// Spawn `f` with `supervisor` as its parent. The supervisor will receive
/// a `Signal` on its registered channel when the child terminates.
pub fn spawn_under(supervisor: Pid, f: impl FnOnce() + Send + 'static) -> JoinHandle {
    let pid = with_sched(|s| {
        let (idx, gen) = s.allocate_slot();
        let pid = Pid::new(idx, gen);
        let stack = Stack::new(ACTOR_STACK_SIZE)
            .expect("stack allocation failed");
        let sp = init_actor_stack(stack.top(), trampoline);
        let slot = &mut s.slots[idx as usize];
        slot.actor = Some(Actor { pid, stack, sp, supervisor });
        slot.state = State::Runnable;
        slot.outstanding_handles = 1;
        slot.outcome = None;
        slot.waiters.clear();
        slot.supervisor_channel = None;
        s.run_queue.push_back(pid);
        pid
    });

    // Stash the closure where `schedule_loop` will find it before the first
    // resume.
    PENDING_CLOSURES.with(|c| {
        c.borrow_mut().push((pid, Box::new(f) as Closure));
    });

    JoinHandle { pid, consumed: false }
}

type Closure = Box<dyn FnOnce() + Send>;

thread_local! {
    /// Closures awaiting their first resume. Keyed by the PID the scheduler
    /// allocated for them in `spawn_under`. The scheduler pops from here in
    /// `pop_pending_closure` right before each first resume.
    static PENDING_CLOSURES: RefCell<Vec<(Pid, Closure)>> = const { RefCell::new(Vec::new()) };
}

fn pop_pending_closure(pid: Pid) -> Option<Closure> {
    PENDING_CLOSURES.with(|c| {
        let mut v = c.borrow_mut();
        v.iter().position(|(p, _)| *p == pid).map(|i| v.swap_remove(i).1)
    })
}

pub fn self_pid() -> Pid {
    current_pid().expect("self_pid() called outside an actor")
}

// ---------------------------------------------------------------------------
// yield_now / park / unpark
// ---------------------------------------------------------------------------

/// Cooperative yield. The current actor goes to the back of the run queue.
pub fn yield_now() {
    // Mark ourselves as needing to be re-queued, then yield.
    YIELD_INTENT.with(|c| c.set(YieldIntent::Yield));
    unsafe { crate::context::switch_to_scheduler() };
}

/// Park the current actor (remove it from the run queue until `unpark`).
pub fn park_current() {
    YIELD_INTENT.with(|c| c.set(YieldIntent::Park));
    unsafe { crate::context::switch_to_scheduler() };
}

/// Wake a parked actor. If the actor isn't parked (already runnable or done)
/// this is a no-op — that's important; channel and join can both fire
/// spurious unparks under some orderings and we want them to be cheap.
/// Also a no-op if the scheduler isn't running (covers channel-sender drop
/// during runtime teardown).
pub fn unpark(pid: Pid) {
    try_with_sched(|s| {
        if let Some(slot) = s.slot_mut(pid) {
            if matches!(slot.state, State::Parked) {
                slot.state = State::Runnable;
                s.run_queue.push_back(pid);
            }
        }
    });
}

/// What an actor wants the scheduler to do when control returns from it.
#[derive(Copy, Clone)]
enum YieldIntent {
    /// Re-queue (yield_now or preemption).
    Yield,
    /// Remove from the run queue (waiting for unpark).
    Park,
}

thread_local! {
    static YIELD_INTENT: std::cell::Cell<YieldIntent> = const { std::cell::Cell::new(YieldIntent::Yield) };
}

// ---------------------------------------------------------------------------
// Supervisor channel registration
// ---------------------------------------------------------------------------

/// Register `sender` as the mailbox for signals about children supervised
/// by `pid`. Idempotent; later calls overwrite.
pub fn register_supervisor_channel(pid: Pid, sender: Sender<Signal>) {
    with_sched(|s| {
        if let Some(slot) = s.slot_mut(pid) {
            slot.supervisor_channel = Some(sender);
        } else {
            panic!("register_supervisor_channel: pid {:?} not found", pid);
        }
    });
}

// ---------------------------------------------------------------------------
// run() — the runtime entry point
// ---------------------------------------------------------------------------

/// Boot the runtime, spawn `initial` as a child of the root supervisor,
/// drive the scheduler until the run queue is empty, tear down.
///
/// The root supervisor is a *sentinel* PID, not a real actor. Signals
/// addressed to it are dropped on the floor — that's what "process exits"
/// means in the spec when nothing escalates further. User code that wants
/// real supervision spawns its own supervisor actor and uses `spawn_under`.
pub fn run<F: FnOnce() + Send + 'static>(initial: F) {
    SCHED.with(|c| {
        assert!(c.borrow().is_none(), "smarm::run() called recursively");
        let mut state = SchedulerState::new();
        state.root_pid = Some(ROOT_PID);
        *c.borrow_mut() = Some(state);
    });

    let initial_handle = spawn(initial);

    schedule_loop();

    // Drop the handle BEFORE the scheduler is torn down — its Drop impl
    // calls `with_sched` to decrement the outstanding-handle count.
    drop(initial_handle);

    // Take the SchedulerState out of the thread-local BEFORE dropping it.
    // Dropping it while still inside SCHED.with's RefCell borrow would
    // re-enter (via channel senders' Drop → unpark → try_with_sched).
    let state = SCHED.with(|c| c.borrow_mut().take());
    drop(state);
    PENDING_CLOSURES.with(|c| c.borrow_mut().clear());
}

/// Reserved sentinel pid for the root supervisor. Never allocated to a
/// real actor; lookups return `None`; signals are dropped.
pub const ROOT_PID: Pid = Pid::new(u32::MAX, u32::MAX);

fn schedule_loop() {
    loop {
        let pid = match with_sched(|s| s.run_queue.pop_front()) {
            Some(p) => p,
            None => return,
        };

        // Look up sp; skip stale or already-reaped pids.
        let sp = match with_sched(|s| {
            s.slot(pid).and_then(|slot| slot.actor.as_ref().map(|a| a.sp))
        }) {
            Some(sp) => sp,
            None => continue,
        };

        // If this is a first resume, move the pending closure to the
        // thread-local the trampoline reads.
        if let Some(b) = pop_pending_closure(pid) {
            set_current_actor_box(b);
        }

        set_actor_sp(sp);
        set_current_pid(pid);
        reset_actor_done();
        YIELD_INTENT.with(|c| c.set(YieldIntent::Yield));

        crate::preempt::reset_timeslice();
        PREEMPTION_ENABLED.with(|c| c.set(true));

        unsafe { switch_to_actor() };

        PREEMPTION_ENABLED.with(|c| c.set(false));
        clear_current_pid();

        let intent = YIELD_INTENT.with(|c| c.get());
        let new_sp = get_actor_sp();

        if is_actor_done() {
            let outcome = take_last_outcome().unwrap_or(Outcome::Exit);
            finalize_actor(pid, outcome);
        } else {
            with_sched(|s| {
                if let Some(slot) = s.slot_mut(pid) {
                    if let Some(actor) = slot.actor.as_mut() {
                        actor.sp = new_sp;
                    }
                    match intent {
                        YieldIntent::Yield => {
                            slot.state = State::Runnable;
                            s.run_queue.push_back(pid);
                        }
                        YieldIntent::Park => {
                            slot.state = State::Parked;
                        }
                    }
                }
            });
        }
    }
}

fn finalize_actor(pid: Pid, outcome: Outcome) {
    // Joiners get the typed Result with the panic payload. The supervisor
    // gets an informational `Signal::Panic` with an empty payload — its job
    // is policy (restart/escalate), not forensics. Users who need the
    // payload in supervision can plumb their own channel.

    let (joiner_outcome, sup_signal) = match outcome {
        Outcome::Exit             => (Outcome::Exit, Signal::Exit(pid)),
        Outcome::Panic(payload)   => (
            Outcome::Panic(payload),
            Signal::Panic(pid, Box::new(()) as Box<dyn std::any::Any + Send>),
        ),
    };

    // Stash outcome, mark Done, collect waiters, drop the actor stack.
    let (waiters, supervisor_pid) = with_sched(|s| {
        let slot = s.slot_mut(pid).expect("finalize_actor: slot vanished");
        let sup = slot.actor.as_ref().map(|a| a.supervisor);
        slot.outcome = Some(joiner_outcome);
        slot.state = State::Done;
        slot.actor = None;
        let w = std::mem::take(&mut slot.waiters);
        (w, sup)
    });

    // Deliver to supervisor (best-effort; ignore SendError).
    if let Some(sup) = supervisor_pid {
        let sender = with_sched(|s| {
            s.slot(sup).and_then(|slot| slot.supervisor_channel.clone())
        });
        if let Some(sender) = sender {
            let _ = sender.send(sup_signal);
        }
    }

    // Unpark joiners.
    for joiner in waiters {
        unpark(joiner);
    }

    // Reclaim if no outstanding handles.
    with_sched(|s| {
        let should_reclaim = match s.slot(pid) {
            Some(slot) => slot.outstanding_handles == 0,
            None => false,
        };
        if should_reclaim {
            reclaim_slot(s, pid);
        }
    });
}