//! Minimal NCCL FFI bindings (hand-written, like the CUDA bindings in
//! xtrain-cuda). Only the collectives data-parallel training needs:
//! unique-id creation, communicator init/destroy, and AllReduce. Mirrors
//! xserv-distributed's FFI.

use std::ffi::c_void;
use std::os::raw::c_char;
use xtrain_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 every rank.
#[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) — DDP all-reduces fp32 gradients.
pub const NCCL_FLOAT32: i32 = 7;

// 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;
    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)
    );
}