//! 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.
//!
//! 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;

/// 16 KiB per thread stack (no guard page! a deep recursion in a thread
/// silently overwrites whatever is below its stack on the heap).
const STACK_SIZE: usize = 16 * 1024;

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

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];

    // 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,
    });

    // 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();
        if scheduler.threads.is_empty() {
            // First spawn: 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 from it.
            scheduler.threads.push(Box::new(Thread {
                name: "main",
                rsp: 0,
                _stack: None,
                kernel_stack_top: 0,
                cr3: Cr3::read().0,
                process: None,
                finished: false,
            }));
        }
        scheduler.threads.push(thread);
    });
}

/// 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;
        // Find the next thread that hasn't finished.
        let next = (1..=n).map(|i| (current + i) % n).find(|&i| !s.threads[i].finished);
        let next = match next {
            Some(next) if next != current => next,
            _ => return, // nobody else to run
        };
        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)
    })
}

/// Sleep for `ticks` timer ticks (~55 ms each), letting other threads run.
pub fn sleep_ticks(ticks: u64) {
    use crate::interrupts::TICKS;
    use core::sync::atomic::Ordering;
    let wake_at = TICKS.load(Ordering::Relaxed) + ticks;
    while TICKS.load(Ordering::Relaxed) < wake_at {
        x86_64::instructions::hlt(); // the timer interrupt may switch threads here
    }
}

// --- 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();
}