//! GDT (Global Descriptor Table) + TSS (Task State Segment).
//!
//! In 64-bit mode segmentation is mostly dead, but the GDT still matters:
//!
//! * **Privilege levels.** The low 2 bits of CS are the *current privilege
//!   level* (CPL): ring 0 = kernel, ring 3 = user. The code segment
//!   descriptors in the GDT say which ring they belong to, so to run user code
//!   we need user code/data segments (DPL 3) next to the kernel ones.
//! * **The TSS.** The GDT is how the CPU finds the TSS, which holds:
//!     - `privilege_stack_table[0]` ("RSP0"): the stack the CPU switches to when
//!       an interrupt/exception arrives while running *ring 3* code. User
//!       stacks can't be trusted, so every thread that runs user code has its
//!       own kernel stack, and the scheduler updates RSP0 on every switch.
//!     - the Interrupt Stack Table (IST): known-good stacks for specific
//!       exceptions. The double-fault handler uses one, so a kernel stack
//!       overflow doesn't turn into a triple fault (= reboot).
//!
//! Segment order matters because of `syscall`/`sysret` (see syscall.rs): they
//! don't read the GDT entries' contents but *compute* selectors from the STAR
//! MSR: kernel CS = base, kernel SS = base + 8; user SS = base2 + 8,
//! user CS = base2 + 16. Hence:
//!
//!   0x00 null | 0x08 kernel code | 0x10 kernel data | 0x18 user data |
//!   0x20 user code | 0x28 TSS (16 bytes, two slots)
//!
//! **One GDT and TSS per CPU.** Each CPU runs a different thread, so each
//! needs its own RSP0 and its own double-fault stack, i.e. its own TSS. (A
//! TSS descriptor is also marked "busy" when loaded, so two CPUs can't share
//! one GDT entry either.) CPU 0 uses the static tables below, which work
//! before the heap exists; the other CPUs get heap-allocated copies from
//! `init_ap`. The selectors are the same on every CPU.

use crate::smp::MAX_CPUS;
use crate::stack::KernelStack;
use alloc::boxed::Box;
use core::ptr::{addr_of, addr_of_mut, null_mut};
use core::sync::atomic::{AtomicPtr, Ordering};
use lazy_static::lazy_static;
use x86_64::VirtAddr;
use x86_64::structures::gdt::{Descriptor, GlobalDescriptorTable, SegmentSelector};
use x86_64::structures::tss::TaskStateSegment;

/// Which IST slot (0..=6) the double-fault handler uses.
pub const DOUBLE_FAULT_IST_INDEX: u16 = 0;

/// CPU 0's TSS. A `static mut` because the scheduler rewrites RSP0 on every
/// context switch (the CPU re-reads it from memory on each ring 3 -> ring 0
/// transition, so writing the field is enough). Only touched with
/// interrupts disabled, by CPU 0 itself.
static mut TSS: TaskStateSegment = TaskStateSegment::new();

/// CPU 0's double-fault stack (static: no heap needed this early). 64 KiB,
/// because after a kernel stack overflow the double-fault handler runs the
/// code that kills the thread on it.
const DOUBLE_FAULT_STACK_SIZE: usize = 64 * 1024;
static mut DOUBLE_FAULT_STACK: [u8; DOUBLE_FAULT_STACK_SIZE] = [0; DOUBLE_FAULT_STACK_SIZE];

/// Each CPU's TSS (CPU 0: the static one, others: leaked boxes).
static TSS_OF: [AtomicPtr<TaskStateSegment>; MAX_CPUS] = [const { AtomicPtr::new(null_mut()) }; MAX_CPUS];

/// Selectors (GDT offsets + requested privilege level) of our segments.
pub struct Selectors {
    pub kernel_code: SegmentSelector,
    pub kernel_data: SegmentSelector,
    pub user_data: SegmentSelector, // RPL 3
    pub user_code: SegmentSelector, // RPL 3
    tss: SegmentSelector,
}

/// A GDT in our standard layout whose TSS entry points at `tss`.
fn build_gdt(tss: &'static TaskStateSegment) -> (GlobalDescriptorTable, Selectors) {
    let mut gdt = GlobalDescriptorTable::new();
    // Order matters for syscall/sysret, see the module comment.
    let kernel_code = gdt.append(Descriptor::kernel_code_segment());
    let kernel_data = gdt.append(Descriptor::kernel_data_segment());
    let user_data = gdt.append(Descriptor::user_data_segment());
    let user_code = gdt.append(Descriptor::user_code_segment());
    let tss = gdt.append(Descriptor::tss_segment(tss));
    (gdt, Selectors { kernel_code, kernel_data, user_data, user_code, tss })
}

lazy_static! {
    // Safety: TSS is a static, so the reference lives forever.
    static ref GDT: (GlobalDescriptorTable, Selectors) = build_gdt(unsafe { &*addr_of!(TSS) });
}

/// The segment selectors, e.g. to build an `iretq` frame into ring 3.
/// (Identical on every CPU.)
pub fn selectors() -> &'static Selectors {
    &GDT.1
}

/// Load `gdt`, reload the segment registers and load the TSS.
fn load(gdt: &'static (GlobalDescriptorTable, Selectors)) {
    use x86_64::instructions::segmentation::{CS, DS, ES, SS, Segment};
    use x86_64::instructions::tables::load_tss;
    gdt.0.load();
    let s = &gdt.1;
    unsafe {
        // The old selectors point into the previous GDT (the bootloader's, or
        // the AP trampoline's), so reload them. (`iretq` checks SS.)
        CS::set_reg(s.kernel_code);
        SS::set_reg(s.kernel_data);
        DS::set_reg(s.kernel_data);
        ES::set_reg(s.kernel_data);
        load_tss(s.tss);
    }
}

/// CPU 0: load the static GDT and TSS (works without a heap).
pub fn init() {
    unsafe {
        // Stacks grow downwards on x86, so hand the CPU the *end* address.
        let start = VirtAddr::from_ptr(addr_of!(DOUBLE_FAULT_STACK));
        (*addr_of_mut!(TSS)).interrupt_stack_table[DOUBLE_FAULT_IST_INDEX as usize] =
            start + DOUBLE_FAULT_STACK_SIZE as u64;
    }
    TSS_OF[0].store(addr_of_mut!(TSS), Ordering::Release);
    load(&GDT);
}

/// Another CPU (`cpu` > 0): build and load its own GDT and TSS, with a
/// guard-paged double-fault stack. These live forever (CPUs never go away).
pub fn init_ap(cpu: usize) {
    let df_stack = KernelStack::new().expect("no memory for a double-fault stack");
    let mut tss = TaskStateSegment::new();
    tss.interrupt_stack_table[DOUBLE_FAULT_IST_INDEX as usize] = VirtAddr::new(df_stack.top());
    core::mem::forget(df_stack); // never freed
    let tss: *mut TaskStateSegment = Box::into_raw(Box::new(tss));
    let gdt = Box::leak(Box::new(build_gdt(unsafe { &*tss })));
    TSS_OF[cpu].store(tss, Ordering::Release);
    load(gdt);
}

/// Set RSP0 of the current CPU: the kernel stack it switches to when an
/// interrupt arrives while user code runs. Called by the scheduler on
/// every context switch (with interrupts off).
pub fn set_kernel_stack(stack_top: u64) {
    let tss = TSS_OF[crate::smp::cpu_id()].load(Ordering::Acquire);
    if !tss.is_null() {
        unsafe { (*tss).privilege_stack_table[0] = VirtAddr::new(stack_top) };
    }
}