//! SMP (symmetric multiprocessing): starting the other CPUs, and per-CPU data.
//!
//! At power-on only one CPU runs: the BSP (bootstrap processor, our CPU 0).
//! The others, the APs (application processors), sit halted, waiting for
//! a wake-up call. To start one:
//!
//!  1. Copy a small piece of 16-bit code, the **trampoline**, to physical
//!     address 0x8000 (below 1 MiB: a starting AP is in 16-bit "real mode",
//!     the mode the original 8086 had, and can only reach the first MiB).
//!  2. Send it an INIT IPI (reset), wait 10+ ms, then a STARTUP IPI with
//!     vector 0x08: the AP starts executing at 0x08 * 4096 = 0x8000.
//!  3. The trampoline switches straight to 64-bit long mode using the BSP's
//!     own settings (page tables, CR0/CR4/EFER) and calls `ap_entry` in
//!     Rust on a fresh stack.
//!  4. `ap_entry` gives the CPU its own GDT/TSS, loads the shared IDT, sets
//!     up `syscall` and its APIC timer, and turns the CPU's current
//!     execution into its *idle thread*. From then on the scheduler hands
//!     it threads like any other CPU.
//!
//! Which CPUs exist? The ACPI MADT table lists them (see acpi.rs).
//!
//! **Which CPU am I?** Every CPU reads the APIC ID register of its own local
//! APIC; a table maps APIC IDs to our CPU numbers 0, 1, 2, ...

use crate::{acpi, apic, gdt, interrupts, memory, stack::KernelStack, thread};
use core::arch::global_asm;
use core::sync::atomic::{AtomicBool, AtomicU8, AtomicU32, AtomicU64, AtomicUsize, Ordering};
use spin::Mutex;
use x86_64::instructions::interrupts::without_interrupts;

/// Most CPUs we handle (bitmasks are u32; 16 keeps the arrays small).
pub const MAX_CPUS: usize = 16;

/// APIC ID -> CPU number.
static APIC_TO_CPU: [AtomicU8; 256] = [const { AtomicU8::new(0) }; 256];
/// CPU number -> APIC ID.
static CPU_TO_APIC: [AtomicU32; MAX_CPUS] = [const { AtomicU32::new(0) }; MAX_CPUS];
/// How many CPUs are running (CPU 0 plus started APs).
static CPU_COUNT: AtomicUsize = AtomicUsize::new(1);
/// CPUs halted in their idle loop (bit per CPU), so new work can wake one.
static IDLE_CPUS: AtomicU32 = AtomicU32::new(0);

/// Per-CPU data the syscall entry reaches through GS (see syscall.rs).
/// The field order is part of the assembly code: offsets 0 and 8.
#[repr(C)]
pub struct PerCpu {
    /// Top of the kernel stack of the thread running on this CPU.
    pub syscall_kernel_rsp: AtomicU64,
    /// Scratch: the user RSP during the first instructions of a syscall.
    pub syscall_user_rsp: AtomicU64,
}
static PER_CPU: [PerCpu; MAX_CPUS] = [const {
    PerCpu { syscall_kernel_rsp: AtomicU64::new(0), syscall_user_rsp: AtomicU64::new(0) }
}; MAX_CPUS];

/// The number (0, 1, ...) of the CPU this code runs on. Only meaningful
/// with interrupts off, or as a hint: the scheduler may move a thread to
/// another CPU at any timer tick.
pub fn cpu_id() -> usize {
    APIC_TO_CPU[apic::id() as usize & 0xFF].load(Ordering::Relaxed) as usize
}

/// The address space (CR3 value) each CPU is running right now.
static ACTIVE_CR3: [AtomicU64; MAX_CPUS] = [const { AtomicU64::new(0) }; MAX_CPUS];
/// Set by a CPU asking CPU n to flush its TLB; CPU n clears it when done.
static FLUSH_PENDING: [AtomicBool; MAX_CPUS] = [const { AtomicBool::new(false) }; MAX_CPUS];

/// The scheduler tells us which address space `cpu` switched to.
pub fn set_active_cr3(cpu: usize, cr3: u64) {
    ACTIVE_CR3[cpu].store(cr3, Ordering::SeqCst);
}

/// **TLB shootdown.** Each CPU caches page-table entries in its TLB. When
/// a page is unmapped, made read-only or replaced (copy-on-write) in an
/// address space that threads are using on *other* CPUs right now, their
/// TLBs may still hold the old entry -- a thread could keep writing to a
/// page that was just freed. So: interrupt every other CPU running the
/// address space `cr3` and wait until each has flushed its TLB. (One
/// thread per process never needs this: no other CPU runs its address
/// space at the same time.) Call with interrupts *on*, after changing the
/// page tables and before freeing any frame they pointed to.
pub fn tlb_shootdown(cr3: u64) {
    let me = without_interrupts(cpu_id);
    let mut targets = 0u32;
    for cpu in 0..cpu_count() {
        if cpu != me && ACTIVE_CR3[cpu].load(Ordering::SeqCst) == cr3 {
            FLUSH_PENDING[cpu].store(true, Ordering::SeqCst);
            let target = CPU_TO_APIC[cpu].load(Ordering::Relaxed);
            without_interrupts(|| apic::send_ipi(target, interrupts::InterruptIndex::TlbShootdown as u8));
            targets |= 1 << cpu;
        }
    }
    // Wait until they're done. Another CPU may be shooting at *us* at the
    // same time while we have interrupts off: answer its request by hand,
    // or we'd wait for each other forever.
    while (0..MAX_CPUS).any(|cpu| targets & (1 << cpu) != 0 && FLUSH_PENDING[cpu].load(Ordering::SeqCst)) {
        without_interrupts(|| {
            if FLUSH_PENDING[cpu_id()].load(Ordering::SeqCst) {
                handle_tlb_shootdown();
            }
        });
        core::hint::spin_loop();
    }
}

/// The shootdown IPI arrived (interrupts.rs): forget all user translations.
pub fn handle_tlb_shootdown() {
    use x86_64::registers::control::Cr3;
    let (frame, flags) = Cr3::read();
    unsafe { Cr3::write(frame, flags) };
    FLUSH_PENDING[cpu_id()].store(false, Ordering::SeqCst);
}

/// How many CPUs are running.
pub fn cpu_count() -> usize {
    CPU_COUNT.load(Ordering::Acquire)
}

/// This CPU's `PerCpu` block.
pub fn this_cpu() -> &'static PerCpu {
    &PER_CPU[cpu_id()]
}

/// Mark this CPU as idle (about to `hlt`) or busy.
pub fn set_idle(cpu: usize, idle: bool) {
    if idle {
        IDLE_CPUS.fetch_or(1 << cpu, Ordering::SeqCst);
    } else {
        IDLE_CPUS.fetch_and(!(1 << cpu), Ordering::SeqCst);
    }
}

/// New work is ready: if another CPU is idle, interrupt it so it picks the
/// work up now rather than at its next timer tick (up to 10 ms later).
pub fn kick_idle_cpu() {
    let idle = IDLE_CPUS.load(Ordering::SeqCst) & !(1 << cpu_id());
    if idle != 0 {
        let cpu = idle.trailing_zeros() as usize;
        let target = CPU_TO_APIC[cpu].load(Ordering::Relaxed);
        without_interrupts(|| apic::send_ipi(target, interrupts::InterruptIndex::Reschedule as u8));
    }
}

// ---------------------------------------------------------------------------
// Starting the other CPUs
// ---------------------------------------------------------------------------

/// Physical address the trampoline is copied to (must be below 1 MiB and
/// 4 KiB-aligned; the STARTUP IPI's vector is this / 4096).
const TRAMPOLINE: u64 = 0x8000;

/// The trampoline's parameter block (filled in by the BSP for each AP).
/// Must match the `.quad`s at `ap_trampoline_params` below.
#[repr(C)]
struct ApParams {
    cr3: u64,
    cr4: u64,
    efer: u64,
    cr0: u64,
    stack_top: u64,
    cpu: u64,
    entry: u64,
}

unsafe extern "C" {
    static ap_trampoline_start: u8;
    static ap_trampoline_params: u8;
    static ap_trampoline_end: u8;
}

// The trampoline. It runs at physical 0x8000 (not where the linker put it),
// so every address is computed as 0x8000 + (label - ap_trampoline_start).
// AT&T syntax, because that's what most 16-bit examples use.
global_asm!(
    r#"
.section .text.ap_trampoline, "ax"
.code16
.global ap_trampoline_start
ap_trampoline_start:
    cli
    cld
    xorw %ax, %ax
    movw %ax, %ds                 # data segment 0: plain physical addresses
    # A temporary GDT with a 64-bit code segment and a data segment.
    lgdtl (ap_gdt_ptr - ap_trampoline_start + 0x8000)
    # Turn on the same CPU features as the BSP: CR4 (PAE: 64-bit page
    # tables, ...), CR3 (the kernel's page tables), EFER (LME = long mode
    # enable, NXE = no-execute bit, SCE = syscall).
    movl (ap_cr4 - ap_trampoline_start + 0x8000), %eax
    movl %eax, %cr4
    movl (ap_cr3 - ap_trampoline_start + 0x8000), %eax
    movl %eax, %cr3
    movl $0xC0000080, %ecx        # EFER's MSR number
    movl (ap_efer - ap_trampoline_start + 0x8000), %eax
    xorl %edx, %edx
    wrmsr
    # CR0: protection + paging at once. With EFER.LME set this goes straight
    # from real mode to long mode (skipping 32-bit protected mode works on
    # QEMU and real CPUs, though the manuals describe the longer route).
    movl (ap_cr0 - ap_trampoline_start + 0x8000), %eax
    movl %eax, %cr0
    # Far jump into the 64-bit code segment (selector 0x08).
    ljmpl $0x08, $(ap_long_mode - ap_trampoline_start + 0x8000)

.code64
ap_long_mode:
    movw $0x10, %ax
    movw %ax, %ds
    movw %ax, %es
    movw %ax, %ss
    xorw %ax, %ax
    movw %ax, %fs
    movw %ax, %gs
    movq (ap_stack - ap_trampoline_start + 0x8000), %rsp
    movq (ap_cpu - ap_trampoline_start + 0x8000), %rdi
    movq (ap_entry - ap_trampoline_start + 0x8000), %rax
    callq *%rax                   # ap_entry(cpu): never returns
1:  hlt
    jmp 1b

.balign 8
ap_gdt:
    .quad 0
    .quad 0x00af9a000000ffff      # 0x08: 64-bit kernel code
    .quad 0x00cf92000000ffff      # 0x10: kernel data
ap_gdt_ptr:
    .word 3 * 8 - 1
    .long ap_gdt - ap_trampoline_start + 0x8000
.balign 8
.global ap_trampoline_params
ap_trampoline_params:
ap_cr3:   .quad 0
ap_cr4:   .quad 0
ap_efer:  .quad 0
ap_cr0:   .quad 0
ap_stack: .quad 0
ap_cpu:   .quad 0
ap_entry: .quad 0
.global ap_trampoline_end
ap_trampoline_end:
.text
"#,
    options(att_syntax)
);

/// The boot stack of the AP being started (it becomes its idle thread's).
static AP_STACK: Mutex<Option<KernelStack>> = Mutex::new(None);
/// Set by a starting AP as soon as it reaches Rust code.
static AP_ALIVE: AtomicBool = AtomicBool::new(false);
/// Set by a starting AP when it's fully set up.
static AP_READY: AtomicBool = AtomicBool::new(false);

/// Bring up multi-core: enable CPU 0's APIC and timer, create CPU 0's idle
/// thread, then start every other CPU listed in the ACPI tables, one by
/// one. Needs the heap and interrupts on. Returns the number of CPUs.
pub fn init() -> usize {
    apic::init_bsp();
    CPU_TO_APIC[0].store(apic::id(), Ordering::Relaxed);
    APIC_TO_CPU[apic::id() as usize].store(0, Ordering::Relaxed);
    thread::init();

    let ids = acpi::cpu_apic_ids().unwrap_or_default();
    copy_trampoline();
    for id in ids {
        let cpu = cpu_count();
        if id == apic::id() || cpu == MAX_CPUS || id > 255 {
            continue;
        }
        APIC_TO_CPU[id as usize].store(cpu as u8, Ordering::Relaxed);
        CPU_TO_APIC[cpu].store(id, Ordering::Relaxed);
        if start_ap(id, cpu) {
            CPU_COUNT.store(cpu + 1, Ordering::Release);
        } else {
            crate::kprintln!("[smp] CPU with APIC id {} did not start", id);
        }
    }
    cpu_count()
}

fn copy_trampoline() {
    let (start, end) = unsafe { (&ap_trampoline_start as *const u8, &ap_trampoline_end as *const u8) };
    let len = end as usize - start as usize;
    assert!(len <= 4096, "AP trampoline too big");
    let dest: *mut u8 = memory::phys_to_virt(x86_64::PhysAddr::new(TRAMPOLINE)).as_mut_ptr();
    unsafe { core::ptr::copy_nonoverlapping(start, dest, len) };
}

/// Start one AP and wait (up to ~1 s) until it reports in.
fn start_ap(apic_id: u32, cpu: usize) -> bool {
    use x86_64::registers::control::{Cr0, Cr4};
    use x86_64::registers::model_specific::Efer;
    let Some(stack) = KernelStack::new() else { return false };
    let cr3 = memory::kernel_l4_frame().start_address().as_u64();
    assert!(cr3 < 1 << 32, "the trampoline loads CR3 with a 32-bit value");
    let params = ApParams {
        cr3,
        cr4: Cr4::read_raw(),
        efer: Efer::read_raw() & !(1 << 10), // not LMA ("long mode active"): the CPU sets that itself
        cr0: Cr0::read_raw(),
        stack_top: stack.top(),
        cpu: cpu as u64,
        entry: ap_entry as *const () as u64,
    };
    let offset = unsafe { &ap_trampoline_params as *const u8 as u64 - &ap_trampoline_start as *const u8 as u64 };
    let dest: *mut ApParams = memory::phys_to_virt(x86_64::PhysAddr::new(TRAMPOLINE + offset)).as_mut_ptr();
    unsafe { dest.write(params) };
    without_interrupts(|| *AP_STACK.lock() = Some(stack));
    AP_ALIVE.store(false, Ordering::SeqCst);
    AP_READY.store(false, Ordering::SeqCst);

    apic::send_init(apic_id);
    wait_ticks(2, || false); // the spec asks for 10 ms; 2 PIT ticks is >= 55 ms
    apic::send_startup(apic_id, (TRAMPOLINE / 4096) as u8);
    // The spec suggests a second STARTUP IPI if the first one got lost.
    if !wait_ticks(1, || AP_ALIVE.load(Ordering::SeqCst)) {
        apic::send_startup(apic_id, (TRAMPOLINE / 4096) as u8);
    }
    wait_ticks(18, || AP_READY.load(Ordering::SeqCst))
}

/// Wait up to `ticks` PIT ticks for `done()`; returns its final value.
fn wait_ticks(ticks: u64, done: impl Fn() -> bool) -> bool {
    let end = interrupts::ticks() + ticks;
    while interrupts::ticks() < end {
        if done() {
            return true;
        }
        core::hint::spin_loop();
    }
    done()
}

/// Where each AP arrives from the trampoline: 64-bit mode, interrupts off,
/// on its boot stack, with the trampoline's temporary GDT.
extern "C" fn ap_entry(cpu: u64) -> ! {
    let cpu = cpu as usize;
    AP_ALIVE.store(true, Ordering::SeqCst);
    gdt::init_ap(cpu);
    interrupts::init_idt();
    crate::syscall::init();
    apic::init_ap();
    let stack = without_interrupts(|| AP_STACK.lock().take()).expect("AP boot stack missing");
    AP_READY.store(true, Ordering::SeqCst);
    thread::run_ap(cpu, stack) // become this CPU's idle thread
}