//! Preemptive kernel threads (beyond the phil-opp series).
//!
//! Each thread has its own stack. The timer interrupt calls `schedule()`,
//! which picks the next thread round-robin and calls `switch_context` to swap
//! stacks. Threads don't have to cooperate: even `loop {}` gets interrupted.
//!
//! How a switch works (all 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 sleeping 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) additionally carry
//! an address space (CR3) and a kernel stack top. On every switch the
//! scheduler loads the next thread's CR3 and points the TSS's RSP0 and the
//! syscall entry at its kernel stack, so interrupts and syscalls from ring 3
//! land on the right stack.
//!
//! Blocking: a thread waiting for something (keyboard input, a child
//! process, a pipe, a timeout) is marked `blocked` and the scheduler skips it
//! until someone calls `wake` (or its sleep deadline passes). That's what
//! `WaitQueue` is for. If *every* thread is blocked, the CPU idles with `hlt`
//! until an interrupt wakes somebody up.
//!
//! Locking rule (single CPU!): any lock that is also taken with interrupts
//! *off* (the scheduler's, a wait queue's, a pipe buffer's...) must *always*
//! be taken with interrupts off. Otherwise a thread could be preempted while
//! holding it, and the next thread would spin on it forever, unable to be
//! preempted itself.
//!
//! A *new* thread has never been interrupted, so `spawn` builds a fake stack
//! that looks like `switch_context` saved it, with the "return address"
//! pointing at `thread_trampoline`, which enables interrupts and calls the
//! thread's function.

use crate::process::Process;
use alloc::{boxed::Box, sync::Arc, vec, vec::Vec};
use core::arch::global_asm;
use spin::Mutex;
use x86_64::instructions::interrupts;
use x86_64::registers::control::Cr3;
use x86_64::structures::paging::PhysFrame;

/// 64 KiB per thread stack. There's no guard page (stacks are plain heap
/// memory), so overflowing one silently overwrites whatever lies below it.
/// We at least *detect* that afterwards: the lowest 8 bytes hold
/// `STACK_CANARY`, checked whenever the scheduler switches away from a
/// thread. (16 KiB was enough until syscalls went through the FAT driver:
/// unoptimized debug builds use big stack frames.)
const STACK_SIZE: usize = 64 * 1024;
const STACK_CANARY: u64 = 0x5AFE_57AC_C0DE_CAFE;

struct Thread {
    name: &'static str,
    /// Saved stack pointer while the thread is not running.
    rsp: u64,
    /// The thread's stack. `None` for the boot thread (uses the bootloader's stack).
    stack: Option<Vec<u8>>,
    /// Top of that stack (0 for the boot thread). Becomes TSS.RSP0 and the
    /// syscall stack while this thread runs.
    kernel_stack_top: u64,
    /// Level-4 page table to run with (the kernel's, or a process's own).
    cr3: PhysFrame,
    /// The process this thread belongs to (None = pure kernel thread).
    process: Option<Arc<Process>>,
    finished: bool,
    /// Waiting for something; the scheduler skips blocked threads.
    blocked: bool,
    /// If blocked by `sleep_ticks`: wake up at this tick.
    wake_at: Option<u64>,
}

struct Scheduler {
    /// Boxed so each Thread has a fixed address even if the Vec grows: we
    /// hand a raw pointer to `rsp` to `switch_context`.
    threads: Vec<Box<Thread>>,
    current: usize,
}

static SCHEDULER: Mutex<Scheduler> = Mutex::new(Scheduler { threads: Vec::new(), current: 0 });

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

/// Start a thread for `process`: runs `entry` (in the kernel) with the
/// process's address space `cr3` active; `entry` then drops to ring 3.
pub fn spawn_with(name: &'static str, entry: fn(), cr3: PhysFrame, process: Option<Arc<Process>>) {
    let mut stack = vec![0u8; STACK_SIZE];
    stack[..8].copy_from_slice(&STACK_CANARY.to_ne_bytes());

    // 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.as_mut_ptr() as u64 + STACK_SIZE as u64) & !0xf; // 16-byte aligned
    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 thread = Box::new(Thread {
        name,
        rsp,
        stack: Some(stack),
        kernel_stack_top: top,
        cr3,
        process,
        finished: false,
        blocked: false,
        wake_at: None,
    });

    // The timer handler also takes this lock, so keep interrupts off while we
    // hold it -- otherwise the handler would spin forever on our lock.
    interrupts::without_interrupts(|| {
        let mut scheduler = SCHEDULER.lock();
        register_main_thread(&mut scheduler);
        // Free the stacks of threads that have finished (we can't do that in
        // exit_current: the thread is still running on its stack there).
        let current = scheduler.current;
        for (i, t) in scheduler.threads.iter_mut().enumerate() {
            if t.finished && i != current {
                t.stack = None;
                t.process = None;
            }
        }
        scheduler.threads.push(thread);
    });
}

/// On first use, register the code that's running right now (the boot/main
/// thread) as thread 0. Its rsp is filled in the first time we switch away.
fn register_main_thread(scheduler: &mut Scheduler) {
    if scheduler.threads.is_empty() {
        scheduler.threads.push(Box::new(Thread {
            name: "main",
            rsp: 0,
            stack: None,
            kernel_stack_top: 0,
            cr3: Cr3::read().0,
            process: None,
            finished: false,
            blocked: false,
            wake_at: None,
        }));
    }
}

/// Called from the timer interrupt (interrupts disabled, EOI already sent).
/// Switches to the next runnable thread, round-robin. Returns when this
/// thread is scheduled again.
pub fn schedule() {
    let (old_rsp, new_rsp) = {
        let mut s = SCHEDULER.lock();
        let n = s.threads.len();
        if n < 2 {
            return; // no threads spawned yet
        }
        let current = s.current;
        // Wake up sleepers whose time has come.
        let now = crate::interrupts::TICKS.load(core::sync::atomic::Ordering::Relaxed);
        for t in s.threads.iter_mut() {
            if t.wake_at.is_some_and(|at| now >= at) {
                t.wake_at = None;
                t.blocked = false;
            }
        }
        // Find the next thread that can run.
        let next = (1..=n)
            .map(|i| (current + i) % n)
            .find(|&i| !s.threads[i].finished && !s.threads[i].blocked);
        let next = match next {
            Some(next) if next != current => next,
            _ => return, // nobody else to run
        };
        if let Some(stack) = &s.threads[current].stack {
            if stack[..8] != STACK_CANARY.to_ne_bytes() {
                panic!("kernel stack overflow in thread '{}'", s.threads[current].name);
            }
        }
        s.current = next;

        // Prepare the CPU for the next thread *before* switching stacks.
        let t = &s.threads[next];
        if t.kernel_stack_top != 0 {
            crate::gdt::set_kernel_stack(t.kernel_stack_top); // interrupts from ring 3
            crate::syscall::set_kernel_stack(t.kernel_stack_top); // `syscall` from ring 3
        }
        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) };
        }

        let old_rsp: *mut u64 = &mut s.threads[current].rsp;
        (old_rsp, s.threads[next].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) };
}

/// 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 s = SCHEDULER.lock();
        s.threads.get(s.current).and_then(|t| t.process.clone())
    })
}

/// End the running thread: mark it finished and switch away for good.
/// (Its stack is kept: we're still running on it until the switch.)
pub fn exit_current() -> ! {
    interrupts::disable();
    {
        let mut s = SCHEDULER.lock();
        let current = s.current;
        s.threads[current].finished = true;
    }
    schedule(); // never returns: finished threads aren't picked again
    crate::hlt_loop(); // only reached if nothing else exists to run
}

/// Name of the running thread (for printing).
pub fn current_name() -> &'static str {
    interrupts::without_interrupts(|| {
        let s = SCHEDULER.lock();
        s.threads.get(s.current).map_or("main", |t| t.name)
    })
}

/// Id of the running thread (its index in the thread list; ids are never reused).
pub fn current_id() -> usize {
    interrupts::without_interrupts(|| {
        let mut s = SCHEDULER.lock();
        register_main_thread(&mut s);
        s.current
    })
}

/// Make thread `id` runnable again (no-op if it isn't blocked).
pub fn wake(id: usize) {
    interrupts::without_interrupts(|| {
        if let Some(t) = SCHEDULER.lock().threads.get_mut(id) {
            t.blocked = false;
            t.wake_at = None;
        }
    });
}

fn is_blocked(id: usize) -> bool {
    SCHEDULER.lock().threads[id].blocked
}

/// Mark the running thread blocked, then run other threads until somebody
/// wakes it. 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.
fn block_current(wake_at: Option<u64>) {
    let me = {
        let mut s = SCHEDULER.lock();
        register_main_thread(&mut s);
        let me = s.current;
        s.threads[me].blocked = true;
        s.threads[me].wake_at = wake_at;
        me
    };
    loop {
        schedule(); // returns once we're picked again -- or if nobody else can run
        if !is_blocked(me) {
            return;
        }
        // Every thread is blocked: idle until an interrupt (timer, keyboard)
        // possibly wakes someone. `enable_and_hlt` = `sti; hlt` with no gap,
        // so the interrupt can't sneak in before the hlt.
        interrupts::enable_and_hlt();
        interrupts::disable();
        if !is_blocked(me) {
            return;
        }
    }
}

/// Sleep for `ticks` timer ticks (~55 ms each). The thread is *blocked*
/// meanwhile, so it costs no CPU time; the scheduler wakes it up.
pub fn sleep_ticks(ticks: u64) {
    use crate::interrupts::TICKS;
    use core::sync::atomic::Ordering;
    let was_enabled = interrupts::are_enabled();
    interrupts::disable();
    let wake_at = TICKS.load(Ordering::Relaxed) + ticks.max(1);
    block_current(Some(wake_at));
    if was_enabled {
        interrupts::enable();
    }
}

/// 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<usize>>,
}

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

    pub fn wait_until<T>(&self, mut condition: impl FnMut() -> Option<T>) -> T {
        let was_enabled = interrupts::are_enabled();
        let result = loop {
            interrupts::disable();
            // Check and go to sleep without interrupts in between: otherwise
            // a wake_all() could happen right after the check, and we'd
            // sleep forever (the classic "lost wake-up" bug).
            if let Some(value) = condition() {
                break value;
            }
            let me = current_id();
            self.waiting.lock().push(me);
            block_current(None);
            // Woken up: loop and check the condition again (somebody else may
            // have consumed what we were waiting for).
        };
        if was_enabled {
            interrupts::enable();
        }
        result
    }

    pub fn wake_all(&self) {
        interrupts::without_interrupts(|| {
            for id in self.waiting.lock().drain(..) {
                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 the timer interrupt, where interrupts are
    // off. 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();
}