//! Per-CPU 数据结构
//!
//! 提供统一的 [`PerCpu<T, A>`] 容器，支持两种使用模式：
//!
//! - **全局定义** — `static` 声明 + 延迟初始化（[`uninit`](PerCpu::uninit) +
//!   [`init`](PerCpu::init)）
//! - **动态定义** — 堆上分配，Drop 时自动释放（[`new`](PerCpu::new) / [`new_in`](PerCpu::new_in)）
//!
//! 数据按 [`Layout::array::<T>(nr_cpus)`](Layout::array) 分配，保证平坦连续布局。
//! 支持自定义分配器（[`PerCpuAlloc`]），默认使用 `#[global_allocator]`。
//!
//! # 示例
//!
//! ```
//! use moa_percpu::PerCpu;
//!
//! // 全局定义
//! static COUNTERS: PerCpu<u64> = PerCpu::uninit();
//!
//! # unsafe {
//! // 初始化（启动阶段调用一次）
//! unsafe { COUNTERS.init(4, || 0) };
//!
//! // 访问
//! assert_eq!(*COUNTERS.cpu(0), 0);
//! unsafe { *COUNTERS.cpu_mut(1) = 42 };
//! assert_eq!(*COUNTERS.cpu(1), 42);
//! # }
//! ```
//!
//! ```
//! use moa_percpu::PerCpu;
//!
//! // 动态定义
//! let pc = PerCpu::new(4, || 0u64);
//! assert_eq!(*pc.cpu(2), 0);
//! // pc 离开作用域时自动释放
//! ```
#![no_std]
#![allow(unsafe_code)]

extern crate alloc;

use core::{
    alloc::Layout,
    cell::UnsafeCell,
    mem::MaybeUninit,
    ptr::{self, NonNull},
    sync::atomic::{AtomicPtr, AtomicUsize, Ordering},
};

/// Per-CPU 分配器接口
///
/// # Safety
///
/// - [`allocate`](PerCpuAlloc::allocate) 返回的指针必须满足 `layout` 的对齐和大小要求
/// - [`deallocate`](PerCpuAlloc::deallocate) 的 `ptr` 和 `layout` 必须与之前
///   [`allocate`](PerCpuAlloc::allocate) 的返回值和参数匹配
pub unsafe trait PerCpuAlloc {
    /// 分配内存
    ///
    /// 返回满足 `layout` 要求的指针。分配失败时应 panic。
    fn allocate(&self, layout: Layout) -> NonNull<u8>;

    /// 释放内存
    ///
    /// # Safety
    ///
    /// `ptr` 必须由同一分配器的 [`allocate`](PerCpuAlloc::allocate)
    /// 返回，`layout` 必须一致。
    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout);
}

/// 默认分配器，转发到 `#[global_allocator]`（ZST，零开销）
pub struct Global;

// Safety: 转发到 Rust global_allocator，满足 Layout 对齐和大小要求。
unsafe impl PerCpuAlloc for Global {
    fn allocate(&self, layout: Layout) -> NonNull<u8> {
        if layout.size() == 0 {
            // Safety: align 始终非零
            return unsafe { NonNull::new_unchecked(layout.align() as *mut u8) };
        }
        // Safety: layout.size() > 0
        let ptr = unsafe { alloc::alloc::alloc(layout) };
        NonNull::new(ptr).expect("allocation failed")
    }

    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
        if layout.size() == 0 {
            return;
        }
        // Safety: 调用者保证 ptr 和 layout 匹配
        unsafe { alloc::alloc::dealloc(ptr.as_ptr(), layout) };
    }
}

/// Per-CPU 数据容器
///
/// `#[repr(C)]` 保证内存布局：`ptr` 在 offset 0，方便汇编直接读取。
///
/// 对于 [`Global`]（ZST），整个结构体仅占 16 字节（`ptr` + `nr_cpus`）。
#[repr(C)]
pub struct PerCpu<T, A: PerCpuAlloc = Global> {
    ptr: AtomicPtr<T>,
    nr_cpus: AtomicUsize,
    // UnsafeCell: init_in 通过 &self 写入此字段
    // #[repr(transparent)]：与 MaybeUninit<A> 布局完全一致
    alloc: UnsafeCell<MaybeUninit<A>>,
}

// Safety: 各 CPU 只通过 cpu_mut 访问自己的 slot，
// 调用者通过禁抢占/中断保证独占。跨 CPU 的 cpu() 共享引用只读。
unsafe impl<T: Send, A: PerCpuAlloc + Send> Send for PerCpu<T, A> {}
unsafe impl<T: Send, A: PerCpuAlloc + Sync> Sync for PerCpu<T, A> {}

/// 分配并初始化平坦数组，返回首元素指针
fn alloc_and_init<T>(
    nr_cpus: usize,
    mut f: impl FnMut() -> T,
    alloc: &impl PerCpuAlloc,
) -> NonNull<T> {
    let layout = Layout::array::<T>(nr_cpus).expect("layout overflow");
    let raw = alloc.allocate(layout).cast::<T>();
    for i in 0..nr_cpus {
        // Safety: raw 指向 nr_cpus 个 T 大小的连续内存
        unsafe { raw.as_ptr().add(i).write(f()) };
    }
    raw
}

impl<T, A: PerCpuAlloc> PerCpu<T, A> {
    /// 创建未初始化的 `PerCpu`，用于 `static` 声明
    ///
    /// # 示例
    ///
    /// ```
    /// use moa_percpu::PerCpu;
    ///
    /// static DATA: PerCpu<u64> = PerCpu::uninit();
    /// ```
    #[inline(always)]
    pub const fn uninit() -> Self {
        Self {
            ptr: AtomicPtr::new(ptr::null_mut()),
            nr_cpus: AtomicUsize::new(0),
            alloc: UnsafeCell::new(MaybeUninit::uninit()),
        }
    }
}

impl<T> PerCpu<T> {
    /// 使用 `#[global_allocator]` 延迟初始化
    ///
    /// # Safety
    ///
    /// 必须恰好调用一次，且在任何并发访问之前完成。
    ///
    /// # Panics
    ///
    /// `nr_cpus` 为 0 时 panic。
    pub unsafe fn init(&self, nr_cpus: usize, f: impl FnMut() -> T) {
        // Safety: 调用者保证单次调用且无并发
        unsafe { self.init_in(nr_cpus, f, Global) };
    }
}

impl<T, A: PerCpuAlloc> PerCpu<T, A> {
    /// 使用自定义分配器延迟初始化
    ///
    /// # Safety
    ///
    /// 必须恰好调用一次，且在任何并发访问之前完成。
    ///
    /// # Panics
    ///
    /// `nr_cpus` 为 0 时 panic。
    pub unsafe fn init_in(&self, nr_cpus: usize, f: impl FnMut() -> T, alloc: A) {
        debug_assert!(self.ptr.load(Ordering::Relaxed).is_null(), "PerCpu already initialized");
        assert!(nr_cpus > 0, "nr_cpus must be > 0");

        let raw = alloc_and_init(nr_cpus, f, &alloc);

        // Safety: UnsafeCell 允许通过 &self 写入；调用者保证 init_in 只调用一次
        unsafe { (*self.alloc.get()).write(alloc) };

        self.nr_cpus.store(nr_cpus, Ordering::Relaxed);
        // Release：确保上面的写入对后续 Acquire 可见
        self.ptr.store(raw.as_ptr(), Ordering::Release);
    }
}

impl<T> PerCpu<T> {
    /// 使用 `#[global_allocator]` 创建 per-CPU 数据
    ///
    /// # Panics
    ///
    /// `nr_cpus` 为 0 时 panic。
    pub fn new(nr_cpus: usize, f: impl FnMut() -> T) -> Self {
        Self::new_in(nr_cpus, f, Global)
    }
}

impl<T, A: PerCpuAlloc> PerCpu<T, A> {
    /// 使用自定义分配器创建 per-CPU 数据
    ///
    /// # Panics
    ///
    /// `nr_cpus` 为 0 时 panic。
    pub fn new_in(nr_cpus: usize, f: impl FnMut() -> T, alloc: A) -> Self {
        assert!(nr_cpus > 0, "nr_cpus must be > 0");

        let raw = alloc_and_init(nr_cpus, f, &alloc);

        Self {
            ptr: AtomicPtr::new(raw.as_ptr()),
            nr_cpus: AtomicUsize::new(nr_cpus),
            alloc: UnsafeCell::new(MaybeUninit::new(alloc)),
        }
    }

    /// CPU 数量
    #[inline(always)]
    pub fn nr_cpus(&self) -> usize {
        self.nr_cpus.load(Ordering::Relaxed)
    }

    #[inline(always)]
    fn data_ptr(&self) -> *mut T {
        let p = self.ptr.load(Ordering::Acquire);
        debug_assert!(!p.is_null(), "PerCpu not initialized");
        p
    }

    /// 获取指定 CPU 的共享引用
    ///
    /// # Panics
    ///
    /// `cpu >= nr_cpus` 时 panic（debug 模式）。
    #[inline(always)]
    pub fn cpu(&self, cpu: usize) -> &T {
        debug_assert!(cpu < self.nr_cpus(), "cpu index out of range");
        // Safety: init/new 后数据有效，cpu 在范围内
        unsafe { &*self.data_ptr().add(cpu) }
    }

    /// 获取指定 CPU 的可变引用
    ///
    /// # Safety
    ///
    /// 调用者必须保证对该 CPU slot 的独占访问（通常：禁抢占 +
    /// 无 IRQ 处理器并发访问同一数据）。
    ///
    /// # Panics
    ///
    /// `cpu >= nr_cpus` 时 panic（debug 模式）。
    #[inline(always)]
    #[allow(clippy::mut_from_ref)]
    pub unsafe fn cpu_mut(&self, cpu: usize) -> &mut T {
        debug_assert!(cpu < self.nr_cpus(), "cpu index out of range");
        // Safety: 调用者保证独占
        unsafe { &mut *self.data_ptr().add(cpu) }
    }

    /// 获取底层裸指针
    ///
    /// 返回堆上平坦数组的首地址，用于汇编等需要直接操作内存的场景。
    #[inline(always)]
    pub fn as_ptr(&self) -> *const T {
        self.data_ptr()
    }
}

impl<T, A: PerCpuAlloc> Drop for PerCpu<T, A> {
    fn drop(&mut self) {
        let raw = *self.ptr.get_mut();
        if raw.is_null() {
            return;
        }
        let nr_cpus = *self.nr_cpus.get_mut();
        for i in 0..nr_cpus {
            // Safety: 每个 slot 在 new_in/init_in 中已初始化
            unsafe { ptr::drop_in_place(raw.add(i)) };
        }
        let layout = Layout::array::<T>(nr_cpus).expect("layout overflow");
        // Safety: alloc 在 new_in/init_in 中已写入
        let alloc = unsafe { self.alloc.get_mut().assume_init_ref() };
        // Safety: ptr 和 layout 匹配
        unsafe { alloc.deallocate(NonNull::new_unchecked(raw.cast()), layout) };
    }
}
