/*************************************************************************
 * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 *
 * See LICENSE.txt for more license information
 *************************************************************************/

#include <sys/types.h>
#include <unistd.h>

#include "mlx5/mlx5dvsymbols.h"

#ifdef NCCL_BUILD_MLX5DV
/* Mlx5dv linking mode. Symbols are pointers to linked MLX5 Direct Verbs */

#define ASSIGN_SYM(container, symbol, name) container->name = &symbol;

ncclResult_t buildMlx5dvSymbols(struct ncclMlx5dvSymbols* mlx5dvSymbols) {
  ASSIGN_SYM(mlx5dvSymbols, mlx5dv_is_supported, mlx5dv_internal_is_supported);
  ASSIGN_SYM(mlx5dvSymbols, mlx5dv_get_data_direct_sysfs_path, mlx5dv_internal_get_data_direct_sysfs_path);
  ASSIGN_SYM(mlx5dvSymbols, mlx5dv_reg_dmabuf_mr, mlx5dv_internal_reg_dmabuf_mr);
  ASSIGN_SYM(mlx5dvSymbols, mlx5dv_query_device, mlx5dv_internal_query_device);
  ASSIGN_SYM(mlx5dvSymbols, mlx5dv_create_qp, mlx5dv_internal_create_qp);
  return ncclSuccess;
}

#else
/* Mlx5dv dynamic loading mode. Symbols are loaded from shared objects. */

#include <dlfcn.h>
#include "core.h"

// MLX5DV Library versioning
#define MLX5DV_VERSION "MLX5_1.8"

ncclResult_t buildMlx5dvSymbols(struct ncclMlx5dvSymbols* mlx5dvSymbols) {
  static void* mlx5dvhandle = NULL;
  void* tmp;
  void** cast;

  mlx5dvhandle = dlopen("libmlx5.so", RTLD_NOW);
  if (!mlx5dvhandle) {
    mlx5dvhandle = dlopen("libmlx5.so.1", RTLD_NOW);
    if (!mlx5dvhandle) {
      INFO(NCCL_INIT, "Failed to open libmlx5.so[.1]");
      goto teardown;
    }
  }

#define LOAD_SYM(handle, symbol, funcptr) \
  do { \
    cast = (void**)&funcptr; \
    tmp = dlvsym(handle, symbol, MLX5DV_VERSION); \
    if (tmp == NULL) { \
      WARN("dlvsym failed on %s - %s version %s", symbol, dlerror(), MLX5DV_VERSION); \
      goto teardown; \
    } \
    *cast = tmp; \
  } while (0)

// Attempt to load a specific symbol version - fail silently
#define LOAD_SYM_VERSION(handle, symbol, funcptr, version) \
  do { \
    cast = (void**)&funcptr; \
    *cast = dlvsym(handle, symbol, version); \
    if (*cast == NULL) { \
      INFO(NCCL_NET, "dlvsym failed on %s - %s version %s", symbol, dlerror(), version); \
    } \
  } while (0)

  LOAD_SYM(mlx5dvhandle, "mlx5dv_is_supported", mlx5dvSymbols->mlx5dv_internal_is_supported);
  // Cherry-pick the mlx5dv_get_data_direct_sysfs_path API from MLX5 1.25
  LOAD_SYM_VERSION(mlx5dvhandle, "mlx5dv_get_data_direct_sysfs_path",
                   mlx5dvSymbols->mlx5dv_internal_get_data_direct_sysfs_path, "MLX5_1.25");
  // Cherry-pick the ibv_reg_dmabuf_mr API from MLX5 1.25
  LOAD_SYM_VERSION(mlx5dvhandle, "mlx5dv_reg_dmabuf_mr", mlx5dvSymbols->mlx5dv_internal_reg_dmabuf_mr, "MLX5_1.25");
  // mlx5dv_query_device is versioned at MLX5_1.0
  LOAD_SYM_VERSION(mlx5dvhandle, "mlx5dv_query_device", mlx5dvSymbols->mlx5dv_internal_query_device, "MLX5_1.0");
  // mlx5dv_create_qp is versioned at MLX5_1.3
  LOAD_SYM_VERSION(mlx5dvhandle, "mlx5dv_create_qp", mlx5dvSymbols->mlx5dv_internal_create_qp, "MLX5_1.3");

  return ncclSuccess;

teardown:
  mlx5dvSymbols->mlx5dv_internal_is_supported = NULL;
  mlx5dvSymbols->mlx5dv_internal_get_data_direct_sysfs_path = NULL;
  mlx5dvSymbols->mlx5dv_internal_reg_dmabuf_mr = NULL;
  mlx5dvSymbols->mlx5dv_internal_query_device = NULL;
  mlx5dvSymbols->mlx5dv_internal_create_qp = NULL;

  if (mlx5dvhandle != NULL) dlclose(mlx5dvhandle);
  return ncclSystemError;
}

#endif
