//! Harts. Each hart has its own kernel stack, and its hart id doubles as
//! its kernel `CoreId`. The kernel is entered under one lock for now (stage
//! 1 of multi-core; the kernel itself stays a sequential program), so
//! userspace runs in parallel and the kernel does not.
//!
//! Harts talk through IPIs carrying two kinds of request:
//! - reschedule: the kernel made work for this hart, or took its thread
//!   away; handled under the lock on the way out of the kernel.
//! - leave root R: a page table this hart may be using is changing or about
//!   to be freed; the hart switches to the kernel's root and flushes its
//!   TLB. This needs no lock, so it is serviced even while a hart spins
//!   waiting for the lock or idles, and the sender (which holds the lock)
//!   can wait for the acknowledgement without deadlocking.

use core::sync::atomic::{AtomicBool, AtomicU32, AtomicUsize, Ordering};

use crate::mmu;

/// Must match linker.ld.
pub const MAX_HARTS: usize = 8;
/// A guard page and a 512 KiB stack (boot builds big tables on it). Must
/// match linker.ld.
pub const STACK_SLOT: usize = 4096 + 0x8_0000;

const SIP_SSIP: usize = 1 << 1;

/// The hart running this code, from which kernel stack it is on.
pub fn hart_id() -> usize {
    let sp: usize;
    // SAFETY: reading sp has no side effects.
    unsafe { core::arch::asm!("mv {0}, sp", out(reg) sp) };
    (sp - stacks_start()) / STACK_SLOT
}

pub fn stacks_start() -> usize {
    unsafe extern "C" {
        static _stacks_start: u8;
    }
    // SAFETY: only the symbol's address is taken.
    unsafe { &_stacks_start as *const u8 as usize }
}

static BOOT_HART: AtomicUsize = AtomicUsize::new(0);
static NUM_HARTS: AtomicUsize = AtomicUsize::new(1);
/// Physical address of the kernel's own root table.
static KERNEL_ROOT: AtomicUsize = AtomicUsize::new(0);

/// The root table each hart translates with.
static ACTIVE_ROOT: [AtomicUsize; MAX_HARTS] = [const { AtomicUsize::new(0) }; MAX_HARTS];
static RESCHEDULE: [AtomicBool; MAX_HARTS] = [const { AtomicBool::new(false) }; MAX_HARTS];
static LEAVE_ROOT: [AtomicUsize; MAX_HARTS] = [const { AtomicUsize::new(0) }; MAX_HARTS];
static LEFT: [AtomicUsize; MAX_HARTS] = [const { AtomicUsize::new(0) }; MAX_HARTS];

pub fn boot_hart() -> usize {
    BOOT_HART.load(Ordering::Relaxed)
}

pub fn num_harts() -> usize {
    NUM_HARTS.load(Ordering::Relaxed)
}

pub fn kernel_root() -> usize {
    KERNEL_ROOT.load(Ordering::Relaxed)
}

/// Records which hart booted and how many there are, and where the
/// kernel's root table is.
pub fn init(boot: usize, harts: usize, kernel_root: usize) {
    BOOT_HART.store(boot, Ordering::Relaxed);
    NUM_HARTS.store(harts, Ordering::Relaxed);
    KERNEL_ROOT.store(kernel_root, Ordering::Relaxed);
}

/// Called by `mmu::activate`.
pub fn note_active_root(root: usize) {
    if let Some(r) = ACTIVE_ROOT.get(hart_id()) {
        r.store(root, Ordering::SeqCst);
    }
}

// ---- the kernel lock ----

static NEXT_TICKET: AtomicU32 = AtomicU32::new(0);
static SERVING: AtomicU32 = AtomicU32::new(0);

/// Takes the kernel lock, first come first served. Leave-root requests
/// are serviced while waiting.
pub fn lock() {
    let ticket = NEXT_TICKET.fetch_add(1, Ordering::Relaxed);
    while SERVING.load(Ordering::Acquire) != ticket {
        service_ipi();
        core::hint::spin_loop();
    }
}

pub fn unlock() {
    SERVING.fetch_add(1, Ordering::Release);
}

// ---- what lock-free fast paths may know ----
//
// Published under the kernel lock just before it is released, read without
// it. A fast path that reads them while another hart holds the lock simply
// takes effect before that hart's operation.

/// Per core: highest waiting priority + 1, 0 if its queue is empty.
static WAITING: [AtomicUsize; MAX_HARTS] = [const { AtomicUsize::new(0) }; MAX_HARTS];
/// Per hart: the running thread's process id + 1 and priority + 1.
static RUNNING_PID: [AtomicUsize; MAX_HARTS] = [const { AtomicUsize::new(0) }; MAX_HARTS];
static RUNNING_PRIORITY: [AtomicUsize; MAX_HARTS] = [const { AtomicUsize::new(0) }; MAX_HARTS];

pub fn publish_waiting(core: usize, highest: Option<u8>) {
    if let Some(w) = WAITING.get(core) {
        w.store(highest.map_or(0, |p| p as usize + 1), Ordering::Release);
    }
}

pub fn publish_running(pid: Option<u16>, priority: Option<u8>) {
    let me = hart_id();
    if let (Some(p), Some(q)) = (RUNNING_PID.get(me), RUNNING_PRIORITY.get(me)) {
        p.store(pid.map_or(0, |x| x as usize + 1), Ordering::Release);
        q.store(priority.map_or(0, |x| x as usize + 1), Ordering::Release);
    }
}

/// The process the thread running on this hart belongs to.
pub fn running_pid() -> Option<usize> {
    let v = RUNNING_PID.get(hart_id())?.load(Ordering::Acquire);
    v.checked_sub(1)
}

/// Whether a Yield by the running thread would change nothing: no thread
/// waits on this core at its priority or above.
pub fn yield_is_noop() -> bool {
    let me = hart_id();
    let (Some(w), Some(p)) = (WAITING.get(me), RUNNING_PRIORITY.get(me)) else {
        return false;
    };
    let mine = p.load(Ordering::Acquire);
    mine != 0 && w.load(Ordering::Acquire) < mine
}

/// Whether this hart has an IPI or a reschedule request it has not acted
/// on; fast paths must then take the slow path.
pub fn has_pending_request() -> bool {
    let sip: usize;
    // SAFETY: reading sip has no side effects.
    unsafe { core::arch::asm!("csrr {0}, sip", out(reg) sip) };
    sip & SIP_SSIP != 0
        || RESCHEDULE
            .get(hart_id())
            .is_some_and(|r| r.load(Ordering::SeqCst))
}

// ---- IPIs ----

/// The lock-free half of an IPI: acknowledges it and, if asked, leaves the
/// root being changed. Safe to call at any time, any number of times.
pub fn service_ipi() {
    let sip: usize;
    // SAFETY: reading sip has no side effects.
    unsafe { core::arch::asm!("csrr {0}, sip", out(reg) sip) };
    if sip & SIP_SSIP == 0 {
        return;
    }
    // Clear first: a request posted after this raises SSIP again.
    // SAFETY: clearing our own pending software interrupt.
    unsafe { core::arch::asm!("csrc sip, {0}", in(reg) SIP_SSIP) };
    let me = hart_id();
    let (Some(leave), Some(left)) = (LEAVE_ROOT.get(me), LEFT.get(me)) else {
        return;
    };
    let root = leave.swap(0, Ordering::SeqCst);
    if root != 0 {
        if ACTIVE_ROOT.get(me).map(|r| r.load(Ordering::SeqCst)) == Some(root) {
            mmu::activate(mmu::page_table::table_at(kernel_root()));
        }
        // SAFETY: flushing translations has no other effect.
        unsafe { core::arch::asm!("sfence.vma zero, zero") };
        left.fetch_add(1, Ordering::SeqCst);
    }
}

/// Asks `hart` to look at its run queue again, on its way out of the
/// kernel.
pub fn kick(hart: usize) {
    if let Some(r) = RESCHEDULE.get(hart) {
        r.store(true, Ordering::SeqCst);
        sbi_send_ipi(hart);
    }
}

/// Consumes this hart's reschedule request.
pub fn take_reschedule() -> bool {
    RESCHEDULE
        .get(hart_id())
        .is_some_and(|r| r.swap(false, Ordering::SeqCst))
}

/// Makes every other hart that translates through `root` leave it and
/// flush, and waits until each has. Called with the kernel lock held, so
/// no hart can start using `root` meanwhile.
pub fn shoot_down(root: usize) {
    let me = hart_id();
    for hart in 0..num_harts() {
        if hart == me || ACTIVE_ROOT[hart].load(Ordering::SeqCst) != root {
            continue;
        }
        let seen = LEFT[hart].load(Ordering::SeqCst);
        LEAVE_ROOT[hart].store(root, Ordering::SeqCst);
        sbi_send_ipi(hart);
        while LEFT[hart].load(Ordering::SeqCst) == seen {
            core::hint::spin_loop();
        }
    }
}

// ---- SBI ----

const SBI_EXT_IPI: usize = 0x735049;
const SBI_EXT_HSM: usize = 0x48534D;

fn sbi_send_ipi(hart: usize) {
    // SAFETY: an SBI call; a0 and a1 come back clobbered.
    unsafe {
        core::arch::asm!(
            "ecall",
            inlateout("a0") 1usize => _,
            inlateout("a1") hart => _,
            in("a6") 0usize,
            in("a7") SBI_EXT_IPI,
        );
    }
}

/// Whether `hart` exists, per SBI HSM `hart_get_status`.
fn sbi_hart_exists(hart: usize) -> bool {
    let error: isize;
    // SAFETY: an SBI call; a0 and a1 come back clobbered.
    unsafe {
        core::arch::asm!(
            "ecall",
            inlateout("a0") hart => error,
            lateout("a1") _,
            in("a6") 2usize,
            in("a7") SBI_EXT_HSM,
        );
    }
    error == 0
}

/// The number of harts, counting up from hart 0 until one is missing.
pub fn probe_harts() -> usize {
    (0..MAX_HARTS)
        .take_while(|h| sbi_hart_exists(*h))
        .count()
        .max(1)
}

/// Starts `hart` at the physical address `entry` with the MMU off.
pub fn sbi_hart_start(hart: usize, entry: usize) -> bool {
    let error: isize;
    // SAFETY: an SBI call; a0 and a1 come back clobbered.
    unsafe {
        core::arch::asm!(
            "ecall",
            inlateout("a0") hart => error,
            inlateout("a1") entry => _,
            in("a2") 0usize,
            in("a6") 0usize,
            in("a7") SBI_EXT_HSM,
        );
    }
    error == 0
}