use crate::{alloc::Heap, arch};
use core::{cell::RefCell, marker::PhantomData, ops::DerefMut, ptr::addr_of_mut};
use static_assertions::{assert_eq_size, assert_not_impl_any};

/// The data that is stored in a per-CPU structure.
#[derive(Debug)]
pub struct CPULocals {
    /// The index of this CPU.
    pub cpu_number: usize,

    /// The heap used by this CPU's allocator.
    pub heap: RefCell<&'static mut Heap>,

    /// A canary for the `CPULocals` being initialized.
    canary: usize,

    // This ensures that the type is not `Send`.
    _phantom: PhantomData<*mut ()>,
}

impl CPULocals {
    /// Creates a new instance of the `CPULocals`, using it to initialize this CPU's `CPULocals`.
    ///
    /// # Safety
    ///
    /// - The CPULocals must not have already been initialized.
    pub unsafe fn init(cpu_number: usize, heap: &'static mut Heap) {
        arch::get_cpu_locals().write(CPULocals {
            cpu_number,
            heap: RefCell::new(heap),
            canary: CANARY,
            _phantom: PhantomData,
        });
    }

    /// Returns the instance of the `CPULocals` for this CPU.
    pub fn get() -> &'static CPULocals {
        let ptr = arch::get_cpu_locals();
        // SAFETY: The entrypoint sets this up for hart0, and we allocate this for other harts.
        unsafe {
            let canary_ptr = addr_of_mut!((*ptr.as_ptr()).canary);
            assert_eq!(
                *canary_ptr, CANARY,
                "CPULocals were not initialized (and we probably just did UB)"
            );
            ptr.as_ref()
        }
    }

    /// Retrieves a reference to the heap.
    ///
    /// # Panics
    ///
    /// - Panics if the CPU-local heap was already borrowed.
    /// - The returned guard will panic on deref if the heap was not initialized.
    pub fn heap(&'static self) -> impl DerefMut<Target = &'static mut Heap> {
        self.heap.borrow_mut()
    }
}

assert_eq_size!(CPULocals, [usize; 4]);
assert_not_impl_any!(CPULocals: Send);

cfg_if::cfg_if! {
    if #[cfg(target_pointer_width = "32")] {
        const CANARY: usize = usize::from_le_bytes(*b"locl");
    } else if #[cfg(target_pointer_width = "64")] {
        const CANARY: usize = usize::from_le_bytes(*b"CPULocal");
    } else {
        compile_error!("unsupported platform");
    }
}