/**
 * @file mvgal_comm.h
 * @brief MVGAL-Comm Collective Communication Library Header
 *
 * MVGAL - Multi-Vendor GPU Aggregation Layer for Linux
 *
 * SPDX-License-Identifier: MIT
 */

#ifndef MVGAL_COMM_H
#define MVGAL_COMM_H

#include <stdint.h>
#include <stddef.h>

#ifdef __cplusplus
extern "C" {
#endif

/* Data types for collective operations */
typedef enum {
    MVGAL_DATATYPE_INT8 = 0,
    MVGAL_DATATYPE_INT16 = 1,
    MVGAL_DATATYPE_INT32 = 2,
    MVGAL_DATATYPE_INT64 = 3,
    MVGAL_DATATYPE_FLOAT = 4,
    MVGAL_DATATYPE_DOUBLE = 5,
    MVGAL_DATATYPE_UINT8 = 6,
    MVGAL_DATATYPE_UINT16 = 7,
    MVGAL_DATATYPE_UINT32 = 8,
    MVGAL_DATATYPE_UINT64 = 9,
} mvgal_datatype_t;

/* Reduction operations */
typedef enum {
    MVGAL_OP_SUM = 0,
    MVGAL_OP_PROD = 1,
    MVGAL_OP_MAX = 2,
    MVGAL_OP_MIN = 3,
    MVGAL_OP_AVG = 4,
} mvgal_op_t;

/**
 * @brief Initialize collective communication context
 * @param num_gpus Number of GPUs
 * @param gpu_ids Array of GPU device IDs
 * @return 0 on success, -1 on failure
 */
int mvgal_comm_init(int num_gpus, const int *gpu_ids);

/**
 * @brief Finalize collective communication context
 */
void mvgal_comm_fini(void);

/**
 * @brief Perform AllReduce operation
 * @param sendbuf Input buffer
 * @param recvbuf Output buffer
 * @param count Number of elements
 * @param datatype Data type
 * @param op Reduction operation
 * @param root_gpu Root GPU for operation
 * @return 0 on success, -1 on failure
 */
int mvgal_allreduce_ring(void *sendbuf, void *recvbuf, size_t count,
                         mvgal_datatype_t datatype, mvgal_op_t op,
                         int root_gpu);

/**
 * @brief Perform AllGather operation
 * @param sendbuf Input buffer (local data)
 * @param recvbuf Output buffer (gathered data)
 * @param count Number of elements to send
 * @param datatype Data type
 * @param root_gpu Root GPU
 * @return 0 on success, -1 on failure
 */
int mvgal_allgather(void *sendbuf, void *recvbuf, size_t count,
                    mvgal_datatype_t datatype, int root_gpu);

/**
 * @brief Perform Broadcast operation
 * @param buffer Data buffer (both input and output)
 * @param count Number of elements
 * @param datatype Data type
 * @param root_gpu Root GPU to broadcast from
 * @return 0 on success, -1 on failure
 */
int mvgal_broadcast(void *buffer, size_t count, mvgal_datatype_t datatype,
                    int root_gpu);

/**
 * @brief Get estimated bandwidth between two GPUs
 * @param src_gpu Source GPU index
 * @param dst_gpu Destination GPU index
 * @param bytes Number of bytes to transfer
 * @return Estimated bandwidth in MB/s
 */
int mvgal_comm_get_bandwidth(int src_gpu, int dst_gpu, size_t bytes);

/**
 * @brief Get estimated latency between two GPUs
 * @param src_gpu Source GPU index
 * @param dst_gpu Destination GPU index
 * @return Estimated latency in microseconds
 */
int mvgal_comm_get_latency(int src_gpu, int dst_gpu);

/* Helper functions */
static inline size_t mvgal_datatype_size(mvgal_datatype_t dt) {
    switch (dt) {
        case MVGAL_DATATYPE_INT8: return 1;
        case MVGAL_DATATYPE_INT16: return 2;
        case MVGAL_DATATYPE_INT32: return 4;
        case MVGAL_DATATYPE_INT64: return 8;
        case MVGAL_DATATYPE_FLOAT: return 4;
        case MVGAL_DATATYPE_DOUBLE: return 8;
        case MVGAL_DATATYPE_UINT8: return 1;
        case MVGAL_DATATYPE_UINT16: return 2;
        case MVGAL_DATATYPE_UINT32: return 4;
        case MVGAL_DATATYPE_UINT64: return 8;
        default: return 4;
    }
}

#ifdef __cplusplus
}
#endif

#endif /* MVGAL_COMM_H */