//! Minimal NCCL FFI bindings (hand-written, like the CUDA bindings).
//! Only the collectives we need for tensor parallelism.

use std::ffi::c_void;
use std::os::raw::c_char;
use xserv_cuda::ffi::CudaStream;

/// Opaque NCCL communicator handle (`ncclComm_t`).
pub type NcclComm = *mut c_void;

/// `ncclUniqueId` is a 128-byte opaque blob shared from rank 0 to all ranks.
#[repr(C)]
#[derive(Clone, Copy)]
pub struct NcclUniqueId {
    pub internal: [c_char; 128],
}

impl Default for NcclUniqueId {
    fn default() -> Self {
        Self { internal: [0; 128] }
    }
}

// ncclDataType_t (subset)
pub const NCCL_FLOAT32: i32 = 7;
pub const NCCL_BF16: i32 = 9;

// ncclRedOp_t
pub const NCCL_SUM: i32 = 0;

// ncclResult_t
pub const NCCL_SUCCESS: i32 = 0;

unsafe extern "C" {
    pub fn ncclGetUniqueId(uid: *mut NcclUniqueId) -> i32;
    // ncclUniqueId is passed BY VALUE (a 128-byte struct) per the NCCL ABI.
    pub fn ncclCommInitRank(comm: *mut NcclComm, nranks: i32, commid: NcclUniqueId, rank: i32) -> i32;
    pub fn ncclCommDestroy(comm: NcclComm) -> i32;
    pub fn ncclAllReduce(
        sendbuff: *const c_void,
        recvbuff: *mut c_void,
        count: usize,
        datatype: i32,
        op: i32,
        comm: NcclComm,
        stream: CudaStream,
    ) -> i32;
    // Point-to-point primitives for pipeline parallelism (Phase 18).
    pub fn ncclSend(
        sendbuff: *const c_void,
        count: usize,
        datatype: i32,
        peer: i32,
        comm: NcclComm,
        stream: CudaStream,
    ) -> i32;
    pub fn ncclRecv(
        recvbuff: *mut c_void,
        count: usize,
        datatype: i32,
        peer: i32,
        comm: NcclComm,
        stream: CudaStream,
    ) -> i32;
    pub fn ncclGroupStart() -> i32;
    pub fn ncclGroupEnd() -> i32;
    pub fn ncclGetErrorString(result: i32) -> *const c_char;
}

pub fn err_string(result: i32) -> String {
    unsafe {
        let p = ncclGetErrorString(result);
        if p.is_null() {
            return format!("nccl error {result}");
        }
        std::ffi::CStr::from_ptr(p).to_string_lossy().into_owned()
    }
}

pub fn check(result: i32, what: &str) {
    assert_eq!(result, NCCL_SUCCESS, "{what} failed: {}", err_string(result));
}