/*
 * Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES
 * SPDX-License-Identifier: BSD-3-Clause
 */

#include "mpi.h"
#include "nccl.h"
#include "cudss.h"
#include <library_types.h>

#include <stdio.h>

extern "C" {

static inline ncclDataType_t cuda_to_nccl_type(cudssDataType_t type) {
    switch(type) {
    case CUDSS_R_32F: return ncclFloat32;
    case CUDSS_R_64F: return ncclFloat64;
    case CUDSS_R_32I: return ncclInt32;
    case CUDSS_R_64I: return ncclInt64;
    //TODO: Proper error
    default: exit(1);
    }
}

static inline size_t get_sizeof_value(cudssDataType_t type) {
    switch(type) {
    case CUDSS_R_32F: return sizeof(float);
    case CUDSS_R_64F: return sizeof(double);
    case CUDSS_R_32I: return 4 * sizeof(char);
    case CUDSS_R_64I: return 8 * sizeof(char);
    //TODO: Proper error
    default: exit(1);
    }
}

static inline ncclRedOp_t cudss_to_nccl_op(cudssOpType_t op) {
    switch(op) {
    case CUDSS_SUM: return ncclSum;
    case CUDSS_MAX: return ncclMax;
    case CUDSS_MIN: return ncclMin;
    //TODO: Proper error
    default: exit(1);
    }
}

static inline MPI_Datatype cuda_to_mpi_type(cudssDataType_t type) {
    switch(type) {
    case CUDSS_R_32F: return MPI_FLOAT;
    case CUDSS_R_64F: return MPI_DOUBLE;
    case CUDSS_R_32I: return MPI_INT;
    case CUDSS_R_64I: return MPI_LONG_LONG_INT;
    //TODO: Proper error
    default: exit(1);
    }
}

static inline MPI_Op cudss_to_mpi_op(cudssOpType_t op) {
    switch(op) {
    case CUDSS_SUM: return MPI_SUM;
    case CUDSS_MAX: return MPI_MAX;
    case CUDSS_MIN: return MPI_MIN;
    //TODO: Proper error
    default: exit(1);
    }
}

/*
 * CUDA-aware MPI
 */

int cudssCommRankDevice(void *comm, int *rank)
{
    return ncclCommUserRank(*((ncclComm_t*)comm), rank);
}

int cudssCommSizeDevice(void *comm, int *size)
{
    return ncclCommCount(*((ncclComm_t*)comm), size);
}

int cudssSendDevice(const void *buffer, int count, cudssDataType_t datatype, int dest,
    int tag, void *comm, cudaStream_t stream)
{
    return ncclSend(buffer, count, cuda_to_nccl_type(datatype), dest,
        *((ncclComm_t*)comm), stream);
}

int cudssRecvDevice(void *buffer, int count, cudssDataType_t datatype, int root,
    int tag, void *comm, cudaStream_t stream)
{
    return ncclRecv(buffer, count, cuda_to_nccl_type(datatype), root,
        *((ncclComm_t*)comm), stream);
}

int cudssBcastDevice(void *buffer, int count, cudssDataType_t datatype, int root,
    void *comm, cudaStream_t stream)
{
    return ncclBcast(buffer, count, cuda_to_nccl_type(datatype), root,
        *((ncclComm_t*)comm), stream);
}

int cudssReduceDevice(const void *sendbuf, void *recvbuf, int count,
    cudssDataType_t datatype, cudssOpType_t op, int root, void *comm,
    cudaStream_t stream)
{
    return ncclReduce(sendbuf, recvbuf, count, cuda_to_nccl_type(datatype),
        cudss_to_nccl_op(op), root, *((ncclComm_t*)comm), stream);
}

int cudssAllreduceDevice(const void *sendbuf, void *recvbuf, int count,
    cudssDataType_t datatype, cudssOpType_t op, void *comm, cudaStream_t stream)
{
    return ncclAllReduce(sendbuf, recvbuf, count, cuda_to_nccl_type(datatype),
        cudss_to_nccl_op(op), *((ncclComm_t*)comm), stream);
}

int cudssScattervDevice(const void *sendbuf, const int *sendcounts,
    const int *displs, cudssDataType_t sendtype, void *recvbuf, int recvcount,
    cudssDataType_t recvtype, int root, void *comm, cudaStream_t stream)
{
    size_t bytes_per_send_elt = get_sizeof_value(sendtype);
    ncclComm_t *nccl_comm = (ncclComm_t*)comm;

    int mpi_size = 1, mpi_rank = 0;
    ncclCommUserRank(*nccl_comm, &mpi_rank);
    ncclCommCount   (*nccl_comm, &mpi_size);

    for (int peer = 0; peer < mpi_size; peer++) {
        if (mpi_rank == root) {
            if (peer == root) {
                cudaMemcpyAsync(recvbuf, (char*)sendbuf + displs[peer] * bytes_per_send_elt,
                    sendcounts[peer] * bytes_per_send_elt, cudaMemcpyDeviceToDevice, stream);
            } else {
                ncclSend((char*)sendbuf + displs[peer] * bytes_per_send_elt,
                    sendcounts[peer], cuda_to_nccl_type(sendtype), peer,
                    *nccl_comm, stream);
            }
        } else if (mpi_rank == peer) {
            ncclRecv(recvbuf, recvcount, cuda_to_nccl_type(recvtype), root,
                *nccl_comm, stream);
            break;
        }
    }
    return 0;
}

int cudssCommSplitDevice(const void *comm, int color, int key, void *newcomm)
{
    return ncclCommSplit(*((ncclComm_t*)comm), color, key, (ncclComm_t*)newcomm,
        NULL/*nccl_config*/);
}

int cudssCommFreeDevice(void *comm)
{
    return ncclCommDestroy(*((ncclComm_t*)comm));
}

/*
 * Host MPI
 */

int cudssCommRankHost(void *comm, int *rank)
{
    return MPI_Comm_rank(*((MPI_Comm*)comm), rank);
}

int cudssCommSizeHost(void *comm, int *size)
{
    return MPI_Comm_size(*((MPI_Comm*)comm), size);
}

int cudssSendHost(const void *buffer, int count, cudssDataType_t datatype, int dest,
    int tag, void *comm, cudaStream_t stream)
{
    return MPI_Send(buffer, count, cuda_to_mpi_type(datatype), dest, tag,
        *((MPI_Comm*)comm));
}

int cudssRecvHost(void *buffer, int count, cudssDataType_t datatype, int root,
    int tag, void *comm, cudaStream_t stream)
{
    return MPI_Recv(buffer, count, cuda_to_mpi_type(datatype), root, tag,
        *((MPI_Comm*)comm), MPI_STATUS_IGNORE);
}

int cudssBcastHost(void *buffer, int count, cudssDataType_t datatype, int root,
    void *comm, cudaStream_t stream)
{
    return MPI_Bcast(buffer, count, cuda_to_mpi_type(datatype), root,
        *((MPI_Comm*)comm));
}

int cudssReduceHost(const void *sendbuf, void *recvbuf, int count,
    cudssDataType_t datatype, cudssOpType_t op, int root, void *comm,
    cudaStream_t stream)
{
    int rank;
    MPI_Comm_rank(*((MPI_Comm*)comm), &rank);
    if (sendbuf == recvbuf && rank == root)
        return MPI_Reduce(MPI_IN_PLACE, recvbuf, count, cuda_to_mpi_type(datatype),
            cudss_to_mpi_op(op), root, *((MPI_Comm*)comm));
    else
        return MPI_Reduce(sendbuf, recvbuf, count, cuda_to_mpi_type(datatype),
            cudss_to_mpi_op(op), root, *((MPI_Comm*)comm));
}

int cudssAllreduceHost(const void *sendbuf, void *recvbuf, int count,
    cudssDataType_t datatype, cudssOpType_t op, void *comm, cudaStream_t stream)
{
    if (sendbuf == recvbuf)
        return MPI_Allreduce(MPI_IN_PLACE, recvbuf, count, cuda_to_mpi_type(datatype),
            cudss_to_mpi_op(op), *((MPI_Comm*)comm));
    else
        return MPI_Allreduce(sendbuf, recvbuf, count, cuda_to_mpi_type(datatype),
            cudss_to_mpi_op(op), *((MPI_Comm*)comm));
}

int cudssScattervHost(const void *sendbuf, const int *sendcounts,
    const int *displs, cudssDataType_t sendtype, void *recvbuf, int recvcount,
    cudssDataType_t recvtype, int root, void *comm, cudaStream_t stream)
{
    int rank;
    MPI_Comm_rank(*((MPI_Comm*)comm), &rank);
    if (sendbuf == recvbuf && rank == root)
        return MPI_Scatterv(sendbuf, sendcounts, displs, cuda_to_mpi_type(sendtype),
            MPI_IN_PLACE, recvcount, cuda_to_mpi_type(recvtype), root,
            *((MPI_Comm*)comm));
    else
        return MPI_Scatterv(sendbuf, sendcounts, displs, cuda_to_mpi_type(sendtype),
            recvbuf, recvcount, cuda_to_mpi_type(recvtype), root,
            *((MPI_Comm*)comm));
}

int cudssCommSplitHost(const void *comm, int color, int key, void *newcomm)
{
    return MPI_Comm_split(*((MPI_Comm*)comm), color, key, (MPI_Comm*)newcomm);
}

int cudssCommFreeHost(void *comm)
{
    return MPI_Comm_free((MPI_Comm*)comm);
}

int cudssDistributedGetProperty(libraryPropertyType propertyType, int *value) {

    if (!value) {
        printf("Invalid output value pointer");
        return 1;
    }

    // These are the publicly available properties.
    switch (propertyType) {
    case MAJOR_VERSION:
        *value = CUDSS_VERSION_MAJOR;
        break;
    case MINOR_VERSION:
        *value = CUDSS_VERSION_MINOR;
        break;
    case PATCH_LEVEL:
        *value = CUDSS_VERSION_PATCH;
        break;
    default:
        printf("Value of propertyType = %d is not supported", static_cast<int>(propertyType));
        return 2;
    }

    return 0;
}

/*
 * Distributed communication service API wrapper binding table (imported by cuDSS).
 * The exposed C symbol must be named as "cudssDistributedInterface".
 */
cudssDistributedInterface_t cudssDistributedInterface = {
    cudssCommRankDevice,
    cudssCommSizeDevice,
    cudssSendDevice,
    cudssRecvDevice,
    cudssBcastDevice,
    cudssReduceDevice,
    cudssAllreduceDevice,
    cudssScattervDevice,
    cudssCommSplitDevice,
    cudssCommFreeDevice,
    cudssCommRankHost,
    cudssCommSizeHost,
    cudssSendHost,
    cudssRecvHost,
    cudssBcastHost,
    cudssReduceHost,
    cudssAllreduceHost,
    cudssScattervHost,
    cudssCommSplitHost,
    cudssCommFreeHost,
    cudssDistributedGetProperty
};

} // extern "C"
