//! Threads and the scheduler: preemptive, multi-core, with sleep/wake.
//!
//! Each thread has its own kernel stack (64 KiB with a guard page, see
//! stack.rs). Every CPU's APIC timer fires ~100 times a second and calls
//! `schedule()`, which puts the running thread at the back of the shared
//! **ready queue** and switches to the thread at the front. Threads don't
//! have to cooperate: even `loop {}` gets interrupted.
//!
//! How a switch works (inside the timer interrupt, interrupts off):
//!
//!   thread A running -> timer IRQ -> CPU pushes A's RIP/RSP/RFLAGS (the
//!   interrupt frame) on A's stack -> the `x86-interrupt` handler saves the
//!   caller-saved registers (rax, rcx, rdx, rsi, rdi, r8-r11) on A's stack ->
//!   `schedule()` -> `switch_context(&A.rsp, B.rsp)`:
//!       push the callee-saved registers (rbx, rbp, r12-r15) on A's stack
//!       save rsp into A.rsp
//!       load rsp from B.rsp            <- now we're on B's stack
//!       pop B's callee-saved registers
//!       ret                            <- returns into B's earlier schedule()
//!   ... which returns into B's timer handler, which restores B's caller-saved
//!   registers and `iretq`s back to wherever B was interrupted.
//!
//! So a waiting thread's whole CPU state is simply sitting on its own stack,
//! and all we have to remember per thread is its stack pointer.
//!
//! Threads that run user code (processes, see process.rs) also carry an
//! address space (CR3). On every switch the scheduler loads the next
//! thread's CR3 and points this CPU's TSS.RSP0 and syscall stack at the
//! thread's kernel stack.
//!
//! **States.** A thread is Ready (in the queue), Running (on some CPU),
//! Blocked (waiting on a `WaitQueue`: keyboard input, a pipe, a child...),
//! Sleeping (until a PIT tick; the PIT handler wakes it) or Finished. Blocked
//! and sleeping threads are not in the queue at all, so they cost no CPU
//! time: no busy-polling anywhere.
//!
//! **Idle threads.** Each CPU has one, run when the queue is empty. It
//! halts the CPU (`hlt`) until an interrupt arrives. A CPU that queues work
//! while another is idle sends it a "reschedule" IPI (smp.rs).
//!
//! **Several CPUs, one queue.** The scheduler's data is behind one spinlock.
//! Two subtle points:
//!   * A thread that was just preempted on CPU A is put in the queue *before*
//!     A has finished switching away from its stack. CPU B must not start
//!     running it on the same stack at the same time! So each thread has an
//!     `on_cpu` flag, cleared only after the switch (`finish_switch`), and
//!     the scheduler skips threads that still have it set.
//!   * Likewise a finished thread's stack (and maybe its process's memory)
//!     can only be freed once no CPU is on it: also in `finish_switch`.
//!
//! **Locking rule.** Any lock that is ever taken with interrupts *off* (the
//! scheduler's, a wait queue's, a pipe buffer's, the frame allocator's...)
//! must *always* be taken with interrupts off. Otherwise a thread could be
//! preempted while holding it, and an interrupt handler on the same CPU
//! that wants it would spin forever: the holder can't run until the
//! handler returns. (Spinning briefly while *another* CPU holds a lock is
//! fine: that CPU keeps running and will release it.)
//!
//! A *new* thread has never been switched away from, so `spawn` builds a
//! fake stack that looks like `switch_context` saved it, with the "return
//! address" pointing at `thread_trampoline`, which calls the thread's function.

use crate::errno::Errno;
use crate::process::{Process, Task};
use crate::smp::{self, MAX_CPUS};
use crate::stack::KernelStack;
use alloc::{
    boxed::Box,
    collections::{BTreeMap, VecDeque},
    string::String,
    sync::Arc,
    vec::Vec,
};
use core::arch::global_asm;
use core::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use spin::Mutex;
use x86_64::instructions::interrupts;
use x86_64::registers::control::Cr3;
use x86_64::structures::paging::PhysFrame;

/// Thread id. Never reused.
pub type Tid = usize;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum State {
    Ready,
    Running,
    Blocked,
    /// Sleeping until this PIT tick.
    Sleeping(u64),
    Finished,
}

struct Thread {
    name: String,
    /// Saved stack pointer while the thread is not running.
    rsp: u64,
    /// The thread's stack. `None` for the boot thread "main", which runs on
    /// the stack the bootloader gave us (that one has a guard page too).
    stack: Option<KernelStack>,
    /// Level-4 page table to run with (the kernel's, or a process's own).
    cr3: PhysFrame,
    /// The user thread (and through it the process) this kernel thread
    /// runs (None = pure kernel thread).
    task: Option<Arc<Task>>,
    state: State,
    /// Some CPU is (still) running on this thread's stack.
    on_cpu: bool,
    /// A wake-up arrived while it was running: its next block returns at once.
    woken: bool,
    /// One of the per-CPU idle threads (never in the ready queue).
    idle: bool,
    /// CPUs that have flushed this thread's stack from their TLB (stack.rs).
    tlb_clean_on: u32,
    /// CPUs this thread has run on, one bit each (for tests and curiosity).
    ran_on: u32,
}

struct Scheduler {
    /// All threads. Boxed so each Thread has a fixed address while the map
    /// changes: `switch_context` gets a raw pointer to `rsp`.
    threads: BTreeMap<Tid, Box<Thread>>,
    /// Threads waiting for a CPU, in order.
    ready: VecDeque<Tid>,
    /// Per CPU: the running thread, the idle thread, and the thread we just
    /// switched away from (for `finish_switch`).
    current: [Option<Tid>; MAX_CPUS],
    idle: [Option<Tid>; MAX_CPUS],
    prev: [Option<Tid>; MAX_CPUS],
    next_id: Tid,
}

static SCHEDULER: Mutex<Scheduler> = Mutex::new(Scheduler {
    threads: BTreeMap::new(),
    ready: VecDeque::new(),
    current: [None; MAX_CPUS],
    idle: [None; MAX_CPUS],
    prev: [None; MAX_CPUS],
    next_id: 0,
});

/// Lock the scheduler. Interrupts must be off (locking rule above).
fn lock() -> spin::MutexGuard<'static, Scheduler> {
    debug_assert!(!interrupts::are_enabled(), "scheduler lock with interrupts on");
    SCHEDULER.lock()
}

/// Per-CPU timer statistics: ticks seen, and ticks that found the CPU busy.
static CPU_TICKS: [AtomicU64; MAX_CPUS] = [const { AtomicU64::new(0) }; MAX_CPUS];
static CPU_BUSY_TICKS: [AtomicU64; MAX_CPUS] = [const { AtomicU64::new(0) }; MAX_CPUS];
/// Is CPU n running its idle thread right now?
static CPU_IDLING: [AtomicBool; MAX_CPUS] = [const { AtomicBool::new(false) }; MAX_CPUS];
/// Kernel stack overflows caught so far.
static STACK_OVERFLOWS: AtomicU64 = AtomicU64::new(0);
/// Woken whenever a thread is gone for good (for `join`).
static THREAD_EXITED: WaitQueue = WaitQueue::new();

impl Scheduler {
    fn thread(&mut self, id: Tid) -> &mut Thread {
        self.threads.get_mut(&id).expect("no such thread")
    }

    fn add(&mut self, thread: Thread) -> Tid {
        let id = self.next_id;
        self.next_id += 1;
        self.threads.insert(id, Box::new(thread));
        id
    }

    /// The thread running on `cpu`. On first use on CPU 0 that's the boot
    /// code, which we adopt as thread "main" (its rsp is saved on its first
    /// switch). APs are registered in `run_ap` before they schedule.
    fn current(&mut self, cpu: usize) -> Tid {
        if let Some(id) = self.current[cpu] {
            return id;
        }
        let mut main = new_thread("main", None, Cr3::read().0, None);
        main.state = State::Running;
        main.on_cpu = true;
        let id = self.add(main);
        self.current[cpu] = Some(id);
        id
    }

    /// Next runnable thread from the queue, skipping threads another CPU is
    /// still switching away from (they stay queued).
    fn pop_ready(&mut self) -> Option<Tid> {
        for _ in 0..self.ready.len() {
            let id = self.ready.pop_front()?;
            match self.threads.get(&id) {
                Some(t) if t.state == State::Ready && !t.on_cpu => return Some(id),
                Some(t) if t.state == State::Ready => self.ready.push_back(id),
                _ => {} // stale entry (can't really happen): drop it
            }
        }
        None
    }

    /// Make a blocked/sleeping thread ready. Returns true if it was queued.
    fn make_ready(&mut self, id: Tid) -> bool {
        let Some(t) = self.threads.get_mut(&id) else { return false };
        match t.state {
            State::Blocked | State::Sleeping(_) => {
                t.state = State::Ready;
                self.ready.push_back(id);
                true
            }
            // Not asleep (yet): it may be just about to block. Make sure it
            // doesn't, or this wake-up would be lost.
            State::Running | State::Ready => {
                t.woken = true;
                false
            }
            State::Finished => false,
        }
    }
}

fn new_thread(name: &str, stack: Option<KernelStack>, cr3: PhysFrame, task: Option<Arc<Task>>) -> Thread {
    Thread {
        name: String::from(name),
        rsp: 0,
        stack,
        cr3,
        task,
        state: State::Ready,
        on_cpu: false,
        woken: false,
        idle: false,
        tlb_clean_on: 0,
        ran_on: 0,
    }
}

/// Create CPU 0's idle thread (smp::init calls this). Before this, a CPU
/// with nothing to run just halts inside whichever thread is blocking.
pub fn init() {
    let id = new_kernel_thread("idle0", idle_entry, crate::memory::kernel_l4_frame(), None, false);
    interrupts::without_interrupts(|| {
        let mut s = lock();
        s.current(0);
        s.thread(id).idle = true;
        s.idle[0] = Some(id);
    });
}

fn idle_entry() {
    idle_loop();
}

/// An AP (from smp.rs) turns its boot code into its idle thread, on the
/// boot stack it was given. Interrupts are off.
pub fn run_ap(cpu: usize, stack: KernelStack) -> ! {
    {
        let mut s = lock();
        let mut idle = new_thread("idle", Some(stack), crate::memory::kernel_l4_frame(), None);
        idle.name = alloc::format!("idle{}", cpu);
        idle.idle = true;
        idle.state = State::Running;
        idle.on_cpu = true;
        idle.tlb_clean_on = 1 << cpu;
        let id = s.add(idle);
        s.current[cpu] = Some(id);
        s.idle[cpu] = Some(id);
    }
    CPU_IDLING[cpu].store(true, Ordering::Relaxed);
    idle_loop()
}

/// What a CPU does when there's nothing to do: run whatever becomes ready,
/// otherwise sleep (`hlt`) until an interrupt.
fn idle_loop() -> ! {
    loop {
        interrupts::disable();
        schedule(); // returns when the queue is empty (again)
        let cpu = smp::cpu_id();
        // Say "I'm idle" *before* the last look at the queue: whoever queues
        // work after that look sees the flag and sends us an IPI.
        smp::set_idle(cpu, true);
        if interrupts::without_interrupts(|| lock().ready.is_empty()) {
            interrupts::enable_and_hlt(); // `sti; hlt` with no gap for a lost IPI
        } else {
            interrupts::enable();
        }
        smp::set_idle(cpu, false);
    }
}

/// Start a new kernel thread running `entry`. Needs the heap and memory.
pub fn spawn(name: &str, entry: fn()) -> Tid {
    spawn_with(name, entry, crate::memory::kernel_l4_frame(), None)
}

/// Start a kernel thread for user thread `task`: runs `entry` (in the
/// kernel) with the process's address space `cr3` active; `entry` then
/// drops to ring 3.
pub fn spawn_with(name: &str, entry: fn(), cr3: PhysFrame, task: Option<Arc<Task>>) -> Tid {
    new_kernel_thread(name, entry, cr3, task, true)
}

fn new_kernel_thread(name: &str, entry: fn(), cr3: PhysFrame, task: Option<Arc<Task>>, queue: bool) -> Tid {
    let stack = KernelStack::new().expect("out of memory for a kernel stack");

    // Build the initial stack, top down (stacks grow towards lower addresses):
    //   [return address = thread_trampoline]
    //   rbp, rbx, r12 (= entry), r13, r14, r15     <- rsp points at r15
    // exactly what `switch_context` pops before its `ret`.
    let top = stack.top();
    let initial: [u64; 7] = [
        0,                                // r15
        0,                                // r14
        0,                                // r13
        entry as *const () as u64,        // r12: picked up by the trampoline
        0,                                // rbx
        0,                                // rbp
        thread_trampoline as *const () as u64, // return address
    ];
    let rsp = top - 8 * initial.len() as u64;
    unsafe { core::ptr::copy_nonoverlapping(initial.as_ptr(), rsp as *mut u64, initial.len()) };

    let mut thread = new_thread(name, Some(stack), cr3, task);
    thread.rsp = rsp;
    let id = interrupts::without_interrupts(|| {
        let mut s = lock();
        s.current(smp::cpu_id()); // adopt "main" if this is the first spawn
        let id = s.add(thread);
        if queue {
            s.ready.push_back(id);
        }
        id
    });
    if queue {
        smp::kick_idle_cpu();
    }
    id
}

/// Switch this CPU to the next thread in the queue, if there is one.
/// Called with interrupts off: from the timer interrupt, `yield_now`, or a
/// thread that just blocked. Returns when this thread runs again (maybe on
/// another CPU!), or right away if there's nothing else to run.
pub fn schedule() {
    let cpu = smp::cpu_id();
    let (old_rsp, new_rsp) = {
        let mut s = lock();
        let Some(cur) = s.current[cpu] else { return }; // no threads yet
        let cur_state = s.thread(cur).state;
        let cur_idle = s.thread(cur).idle;
        let next = match s.pop_ready() {
            Some(next) => next,
            // Nobody waiting: a running thread just carries on...
            None if cur_state == State::Running => return,
            // ...one that blocked or finished hands over to the idle thread.
            None => match s.idle[cpu] {
                Some(idle) if idle != cur => idle,
                _ => return, // no idle thread yet: the caller halts instead
            },
        };
        // Back of the line for the current thread (if it can still run).
        if cur_state == State::Running {
            if cur_idle {
                s.thread(cur).state = State::Ready; // idle threads aren't queued
            } else {
                s.thread(cur).state = State::Ready;
                s.ready.push_back(cur);
            }
        }

        let bit = 1u32 << cpu;
        let t = s.thread(next);
        t.state = State::Running;
        t.on_cpu = true;
        t.ran_on |= bit;
        if let Some(task) = &t.task {
            task.process.ran_on.fetch_or(bit, core::sync::atomic::Ordering::Relaxed);
            // Thread-local storage: each user thread has its own FS base
            // (set with arch_prctl). The CPU has one FS_BASE register, so
            // it's reloaded on every switch to a user thread.
            x86_64::registers::model_specific::FsBase::write(x86_64::VirtAddr::new_truncate(
                task.fs_base.load(core::sync::atomic::Ordering::Relaxed),
            ));
        }
        if let Some(stack) = &t.stack {
            if t.tlb_clean_on & bit == 0 {
                stack.flush_local_tlb(); // see "Reusing a slot" in stack.rs
                t.tlb_clean_on |= bit;
            }
            // Where interrupts (TSS.RSP0) and `syscall` from ring 3 land.
            crate::gdt::set_kernel_stack(stack.top());
            crate::syscall::set_kernel_stack(stack.top());
        }
        if Cr3::read().0 != t.cr3 {
            // Switch address space. Safe for the code running right now:
            // all kernel mappings (incl. every kernel stack) are identical in
            // every address space. Writing CR3 also flushes the TLB.
            unsafe { Cr3::write(t.cr3, Cr3::read().1) };
        }
        // Which address space this CPU runs now (TLB shootdowns, memory.rs).
        smp::set_active_cr3(cpu, t.cr3.start_address().as_u64());
        CPU_IDLING[cpu].store(t.idle, Ordering::Relaxed);
        let new_rsp = t.rsp;
        s.current[cpu] = Some(next);
        s.prev[cpu] = Some(cur);
        let old_rsp: *mut u64 = &mut s.thread(cur).rsp;
        (old_rsp, new_rsp)
        // The lock is released here, *before* switching: the next thread
        // will never come back to release a lock we still held.
    };
    unsafe { switch_context(old_rsp, new_rsp) };
    finish_switch();
}

/// Right after a switch, on the new thread: the previous thread's stack is
/// no longer in use. Clear its `on_cpu` so other CPUs may run it, and if it
/// had finished, remove it and free its stack (and, if it was the last user
/// of its process, the process's memory, via `Drop for Process`).
fn finish_switch() {
    let cpu = smp::cpu_id();
    let (leftovers, requeued) = {
        let mut s = lock();
        let Some(prev) = s.prev[cpu].take() else { return };
        let t = s.thread(prev);
        t.on_cpu = false;
        let state = t.state;
        // A finished thread: take out what it owns (its stack, and its
        // share of the process) to free below, outside the lock.
        let leftovers = (state == State::Finished).then(|| (prev, t.stack.take(), t.task.take()));
        (leftovers, state == State::Ready)
    };
    if requeued {
        smp::kick_idle_cpu(); // it may have been skipped while on_cpu was set
    }
    if let Some((dead, stack, process)) = leftovers {
        // Free first, *then* remove it from the table: `join` waits for the
        // removal, so once `join` returns the memory really is free (the
        // leak tests count frames right after joining).
        drop(stack); // with interrupts still off: the allocators allow that
        drop(process);
        lock().threads.remove(&dead);
        THREAD_EXITED.wake_all();
    }
}

/// Every CPU's APIC timer tick: count it for the statistics, then preempt.
pub fn timer_tick() {
    let cpu = smp::cpu_id();
    CPU_TICKS[cpu].fetch_add(1, Ordering::Relaxed);
    if !CPU_IDLING[cpu].load(Ordering::Relaxed) {
        CPU_BUSY_TICKS[cpu].fetch_add(1, Ordering::Relaxed);
    }
    schedule();
}

/// (ticks, busy ticks) of CPU `cpu`'s APIC timer so far.
pub fn cpu_ticks(cpu: usize) -> (u64, u64) {
    (CPU_TICKS[cpu].load(Ordering::Relaxed), CPU_BUSY_TICKS[cpu].load(Ordering::Relaxed))
}

/// Give up the rest of this time slice (also callable outside interrupts).
pub fn yield_now() {
    interrupts::without_interrupts(schedule);
}

/// The process the running thread belongs to, if any.
pub fn current_process() -> Option<Arc<Process>> {
    interrupts::without_interrupts(|| {
        let mut s = lock();
        let cur = s.current(smp::cpu_id());
        s.thread(cur).task.as_ref().map(|t| t.process.clone())
    })
}

/// The user thread the running kernel thread carries, if any.
pub fn current_task() -> Option<Arc<Task>> {
    interrupts::without_interrupts(|| {
        let mut s = lock();
        let cur = s.current(smp::cpu_id());
        s.thread(cur).task.clone()
    })
}

/// End the running thread: mark it finished and switch away for good. The
/// next thread on this CPU frees its stack (see `finish_switch`).
pub fn exit_current() -> ! {
    interrupts::disable();
    {
        let mut s = lock();
        let cur = s.current(smp::cpu_id());
        s.thread(cur).state = State::Finished;
    }
    loop {
        schedule(); // never returns once there's an idle thread
        // No idle thread (scheduler not fully started): wait for anything.
        interrupts::enable_and_hlt();
        interrupts::disable();
    }
}

/// Name of the running thread (for printing).
pub fn current_name() -> String {
    interrupts::without_interrupts(|| {
        let mut s = lock();
        let cur = s.current(smp::cpu_id());
        s.thread(cur).name.clone()
    })
}

/// Id of the running thread.
pub fn current_id() -> Tid {
    interrupts::without_interrupts(|| lock().current(smp::cpu_id()))
}

/// Is the scheduler fully running (CPU 0 has its idle thread)?
pub fn scheduler_started() -> bool {
    interrupts::without_interrupts(|| lock().idle[0].is_some())
}

/// Does thread `id` still exist (not finished, or finished but its stack
/// not freed yet)?
pub fn exists(id: Tid) -> bool {
    interrupts::without_interrupts(|| lock().threads.contains_key(&id))
}

/// Wait until thread `id` is completely gone (finished and freed).
pub fn join(id: Tid) {
    THREAD_EXITED.wait_until(|| (!lock().threads.contains_key(&id)).then_some(()));
}

/// Number of threads (including idle threads and "main").
pub fn thread_count() -> usize {
    interrupts::without_interrupts(|| lock().threads.len())
}

/// CPUs thread `id` has run on so far (bit n = CPU n), if it exists.
pub fn ran_on(id: Tid) -> Option<u32> {
    interrupts::without_interrupts(|| lock().threads.get(&id).map(|t| t.ran_on))
}

/// How many kernel stack overflows have been caught.
pub fn stack_overflows() -> u64 {
    STACK_OVERFLOWS.load(Ordering::Relaxed)
}

/// Make thread `id` runnable again (if it's blocked or sleeping), and poke
/// an idle CPU to run it.
pub fn wake(id: Tid) {
    if interrupts::without_interrupts(|| lock().make_ready(id)) {
        smp::kick_idle_cpu();
    }
}

/// PIT tick (CPU 0's timer interrupt): wake every thread whose sleep is over.
pub fn wake_sleepers(now: u64) {
    let woke = {
        let mut s = lock();
        let due: Vec<Tid> = s
            .threads
            .iter()
            .filter(|(_, t)| matches!(t.state, State::Sleeping(at) if at <= now))
            .map(|(id, _)| *id)
            .collect();
        for &id in &due {
            s.make_ready(id);
        }
        !due.is_empty()
    };
    if woke {
        smp::kick_idle_cpu();
    }
}

/// Block the running thread (until `wake`, or until PIT tick `until`), and
/// run other threads meanwhile. Must be called with interrupts *off*, so a
/// wake-up can't slip in between the caller checking its condition and us
/// blocking; returns with interrupts still off. Returns at once if a
/// wake-up already arrived (`woken`).
fn block_current(until: Option<u64>) {
    let me = {
        let mut s = lock();
        let me = s.current(smp::cpu_id());
        let t = s.thread(me);
        if t.woken {
            t.woken = false;
            return;
        }
        t.state = match until {
            Some(at) => State::Sleeping(at),
            None => State::Blocked,
        };
        me
    };
    loop {
        schedule(); // returns once we've been woken and picked again
        if done_waiting(me) {
            return;
        }
        // We're still here: there was nothing (not even an idle thread) to
        // switch to. Halt until an interrupt possibly wakes us.
        // `enable_and_hlt` = `sti; hlt` with no gap for the interrupt.
        interrupts::enable_and_hlt();
        interrupts::disable();
        if done_waiting(me) {
            return;
        }
    }
}

/// After `schedule` returned to a thread that blocked: has it been woken?
/// If it was woken but never left the CPU, it's Ready and queued: take it
/// back out of the queue, it's running right here.
fn done_waiting(me: Tid) -> bool {
    let mut s = lock();
    match s.thread(me).state {
        State::Blocked | State::Sleeping(_) => false,
        State::Ready => {
            s.thread(me).state = State::Running;
            s.ready.retain(|&id| id != me);
            true
        }
        _ => true,
    }
}

/// Clear a stale wake-up flag before starting to wait for something new.
fn clear_woken() {
    let mut s = lock();
    let me = s.current(smp::cpu_id());
    s.thread(me).woken = false;
}

/// Sleep for `ticks` PIT ticks (~55 ms each). The thread is off the run
/// queue meanwhile, so it costs no CPU time; the PIT handler wakes it.
pub fn sleep_ticks(ticks: u64) {
    let was_enabled = interrupts::are_enabled();
    interrupts::disable();
    let wake_at = crate::interrupts::ticks() + ticks.max(1);
    clear_woken();
    while crate::interrupts::ticks() < wake_at {
        block_current(Some(wake_at)); // may return early on a stray wake-up
    }
    if was_enabled {
        interrupts::enable();
    }
}

/// `sleep_ticks` that a signal can cut short: `Err(ERESTARTNOHAND)` then
/// (nanosleep and pause are never restarted after a handler ran).
pub fn sleep_ticks_intr(ticks: u64) -> Result<(), Errno> {
    let was_enabled = interrupts::are_enabled();
    interrupts::disable();
    let wake_at = crate::interrupts::ticks() + ticks.max(1);
    clear_woken();
    let mut result = Ok(());
    while crate::interrupts::ticks() < wake_at {
        if crate::signal::interrupted() {
            result = Err(Errno::ERESTARTNOHAND);
            break;
        }
        block_current(Some(wake_at));
    }
    if was_enabled {
        interrupts::enable();
    }
    result
}

/// Block until a signal arrives (pause()).
pub fn pause() -> Errno {
    static NEVER: WaitQueue = WaitQueue::new();
    match NEVER.wait_until_intr(|| None::<()>) {
        Err(e) => if e == Errno::ERESTARTSYS { Errno::ERESTARTNOHAND } else { e },
        Ok(()) => unreachable!(),
    }
}

/// The running thread overflowed its kernel stack (the page-fault or
/// double-fault handler saw an access to a guard page). Say so, and kill
/// it: its process if it has one, else just the thread.
pub fn die_of_stack_overflow(addr: u64) -> ! {
    STACK_OVERFLOWS.fetch_add(1, Ordering::Relaxed);
    let name = current_name();
    crate::kprintln!(
        "[kernel] kernel stack overflow in thread '{}' (hit guard page {:#x}); killing it",
        name,
        addr & !0xfff
    );
    if current_process().is_some() {
        crate::process::exit_group(crate::process::ExitStatus::Killed(crate::signal::SIGSEGV, "kernel stack overflow"));
    }
    exit_current()
}

/// Threads waiting for some event (like a Linux `wait_queue_head_t`).
///
/// Usage: `queue.wait_until(|| condition)` blocks until the closure returns
/// `Some(value)`; whoever makes the condition true calls `queue.wake_all()`.
/// The closure runs with interrupts off, so it must be quick and only touch
/// locks that follow the locking rule at the top of this file.
pub struct WaitQueue {
    waiting: Mutex<Vec<Tid>>,
}

impl WaitQueue {
    pub const fn new() -> Self {
        WaitQueue { waiting: Mutex::new(Vec::new()) }
    }

    pub fn wait_until<T>(&self, condition: impl FnMut() -> Option<T>) -> T {
        match self.wait(None, false, condition) {
            Ok(Some(value)) => value,
            _ => unreachable!("no deadline, not interruptible"),
        }
    }

    /// Like `wait_until`, but give up at PIT tick `deadline`: returns
    /// `None` if the condition still wasn't true by then. (Network code uses
    /// this for timeouts: connect, DNS, ping, `SO_RCVTIMEO`.)
    pub fn wait_until_deadline<T>(&self, deadline: u64, condition: impl FnMut() -> Option<T>) -> Option<T> {
        self.wait(Some(deadline), false, condition).ok().flatten()
    }

    /// Like `wait_until`, but a signal for the process interrupts the wait:
    /// `Err(ERESTARTSYS)`. Every wait a program can get stuck in (reading
    /// the keyboard or a pipe, wait4, sockets, ...) uses this, or Ctrl+C
    /// couldn't stop a program blocked in `read`. On the way back to user
    /// space the signal code turns ERESTARTSYS into EINTR or restarts the
    /// call (see signal.rs).
    pub fn wait_until_intr<T>(&self, condition: impl FnMut() -> Option<T>) -> Result<T, Errno> {
        self.wait(None, true, condition).map(|v| v.expect("no deadline"))
    }

    /// Interruptible *and* with a deadline (`Ok(None)` = timed out).
    pub fn wait_until_deadline_intr<T>(&self, deadline: u64, condition: impl FnMut() -> Option<T>) -> Result<Option<T>, Errno> {
        self.wait(Some(deadline), true, condition)
    }

    fn wait<T>(&self, deadline: Option<u64>, interruptible: bool, mut condition: impl FnMut() -> Option<T>) -> Result<Option<T>, Errno> {
        let was_enabled = interrupts::are_enabled();
        let me = current_id();
        let result = loop {
            interrupts::disable();
            // The order matters (the classic "lost wake-up" bug): put
            // ourselves on the list *first*, then check. If the condition
            // becomes true on another CPU after our check, its wake_all()
            // finds us on the list and sets `woken`, so block_current
            // returns straight away instead of sleeping forever. (Signals
            // wake us with `wake(tid)` directly, which works the same way.)
            clear_woken();
            self.waiting.lock().push(me);
            if let Some(value) = condition() {
                break Ok(Some(value));
            }
            if interruptible && crate::signal::interrupted() {
                break Err(Errno::ERESTARTSYS);
            }
            if deadline.is_some_and(|d| crate::interrupts::ticks() >= d) {
                break Ok(None);
            }
            // Sleep until woken (or the deadline tick, whichever is first).
            block_current(deadline);
            // Woken up: take us off the list (a timer or signal wake-up
            // leaves us on it), then check again -- somebody else may have
            // consumed what we were waiting for.
            self.waiting.lock().retain(|&id| id != me);
        };
        self.waiting.lock().retain(|&id| id != me);
        if was_enabled {
            interrupts::enable();
        }
        result
    }

    pub fn wake_all(&self) {
        let waiting: Vec<Tid> = interrupts::without_interrupts(|| self.waiting.lock().drain(..).collect());
        for id in waiting {
            wake(id);
        }
    }
}

// --- assembly ---------------------------------------------------------------

unsafe extern "C" {
    /// Save callee-saved registers + rsp into `*old_rsp`, load `new_rsp`,
    /// restore that thread's registers and return into it.
    fn switch_context(old_rsp: *mut u64, new_rsp: u64);
    /// First "return address" of every new thread.
    fn thread_trampoline();
}

// System V calling convention: 1st argument in rdi, 2nd in rsi. A normal
// function call already saved the caller-saved registers, so only the
// callee-saved ones (rbx, rbp, r12-r15) need saving here.
global_asm!(
    ".global switch_context",
    "switch_context:",
    "    push rbp",
    "    push rbx",
    "    push r12",
    "    push r13",
    "    push r14",
    "    push r15",
    "    mov [rdi], rsp", // old thread's rsp -> *old_rsp
    "    mov rsp, rsi",   // switch stacks!
    "    pop r15",
    "    pop r14",
    "    pop r13",
    "    pop r12",
    "    pop rbx",
    "    pop rbp",
    "    ret",            // pops the new thread's return address
    "",
    ".global thread_trampoline",
    "thread_trampoline:",
    "    mov rdi, r12",   // entry function pointer -> 1st argument
    "    and rsp, -16",   // ABI: stack 16-byte aligned before `call`
    "    call {start}",
    "    ud2",            // thread_start never returns
    start = sym thread_start,
);

/// Rust side of a new thread's life.
extern "C" fn thread_start(entry: usize) -> ! {
    // We arrived here from inside `schedule`, like any switched-to thread:
    // do the same clean-up as `schedule` does after its switch.
    finish_switch();
    // Interrupts are off in `schedule`. Turn them back on, or this thread
    // could never be preempted.
    interrupts::enable();

    let entry: fn() = unsafe { core::mem::transmute(entry) };
    entry();

    // The thread function returned: mark it finished so the scheduler skips it.
    exit_current();
}