#include <c10/cuda/PeerToPeerAccess.h>

#include <c10/cuda/CUDACachingAllocator.h>
#include <c10/cuda/CUDAException.h>
#include <c10/cuda/CUDAGuard.h>
#if !defined(USE_ROCM) && defined(PYTORCH_C10_DRIVER_API_SUPPORTED)
#include <c10/cuda/driver_api.h>
#endif
#include <c10/util/Exception.h>
#include <c10/util/Logging.h>
#include <c10/util/irange.h>

#include <iomanip>
#include <sstream>
#include <vector>

namespace c10::cuda {

static std::vector<int8_t> p2pAccessEnabled_;
static std::vector<int8_t> fabricAccessEnabled_;
static std::vector<int> fabricCliqueId_;
static int64_t num_devices_ = -1;

namespace detail {

void init_p2p_access_cache(int64_t num_devices) {
  // p2pAccessEnabled records if p2p copies are allowed between pairs of
  // devices. Values include "1" (copy allowed), "0" (copy not allowed), and
  // "-1" (unknown).
  // Currently the max number of gpus in P2P group is 8, so if there are more
  // we enable P2P in groups of 8
  p2pAccessEnabled_.clear();
  p2pAccessEnabled_.resize(num_devices * num_devices, -1);
  num_devices_ = num_devices;

  for (const auto i : c10::irange(num_devices)) {
    p2pAccessEnabled_[i * num_devices + i] = 1;
  }
  fabricAccessEnabled_.clear();
  fabricAccessEnabled_.resize(num_devices, -1);
  fabricCliqueId_.clear();
  fabricCliqueId_.resize(num_devices, kCliqueIdNotQueried);
}

} // namespace detail

bool get_p2p_access(c10::DeviceIndex dev, c10::DeviceIndex dev_to_access) {
  TORCH_CHECK(
      num_devices_ >= 0,
      "p2p access cache not initialized. "
      "Ensure c10::cuda::detail::init_p2p_access_cache() is called first.");
  TORCH_CHECK(
      dev >= 0 && dev < num_devices_,
      static_cast<int>(dev),
      " is not a valid device");
  TORCH_CHECK(
      dev_to_access >= 0 && dev_to_access < num_devices_,
      static_cast<int>(dev_to_access),
      " is not a valid device");

  auto& cache = p2pAccessEnabled_[dev * num_devices_ + dev_to_access];

  if (cache != -1) {
    return cache;
  }

  int result = 0;
  C10_CUDA_CHECK(cudaDeviceCanAccessPeer(&result, dev, dev_to_access));
  cache = result ? 1 : 0;
  if (cache) {
    CUDACachingAllocator::enablePeerAccess(dev, dev_to_access);
  }

  return cache;
}

namespace {
#if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12040 && \
    defined(PYTORCH_C10_DRIVER_API_SUPPORTED)

nvmlDevice_t get_nvml_device(c10::DeviceIndex dev) {
  static bool nvml_init [[maybe_unused]] = []() {
    TORCH_INTERNAL_ASSERT(NVML_SUCCESS == DriverAPI::get()->nvmlInit_v2_());
    return true;
  }();

  // Direct CUDA call instead of at::cuda::getDeviceProperties()
  cudaDeviceProp prop{};
  C10_CUDA_CHECK(cudaGetDeviceProperties(&prop, dev));

  char
      pci_id // NOLINT(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
          [NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE];
  snprintf(
      pci_id,
      sizeof(pci_id),
      NVML_DEVICE_PCI_BUS_ID_FMT,
      prop.pciDomainID,
      prop.pciBusID,
      prop.pciDeviceID);

  nvmlDevice_t nvml_device = nullptr;
  TORCH_INTERNAL_ASSERT(
      NVML_SUCCESS ==
      DriverAPI::get()->nvmlDeviceGetHandleByPciBusId_v2_(
          pci_id, &nvml_device));
  return nvml_device;
}

bool isFabricSupported() {
  // 1. try allocating memory with FABRIC handle type
  CUmemGenericAllocationHandle handle = 0;
  CUmemAllocationProp prop = {};
  prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
  prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC;
  prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;

  size_t granularity{};
  const auto driver_api = DriverAPI::get();
  C10_CUDA_DRIVER_CHECK(driver_api->cuMemGetAllocationGranularity_(
      &granularity, &prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED));

  auto status = driver_api->cuMemCreate_(&handle, granularity, &prop, 0);
  if (status != CUDA_SUCCESS) {
    LOG(INFO)
        << "status " << status
        << " Could not allocate memory with FABRIC handle, falling back to fd handle exchange\n";
    return false;
  }

  // 2. check export
  CUmemFabricHandle sharedHandle;
  status = driver_api->cuMemExportToShareableHandle_(
      &sharedHandle, handle, CU_MEM_HANDLE_TYPE_FABRIC, 0);
  if (status != CUDA_SUCCESS) {
    LOG(INFO)
        << "status " << status
        << " Could not export FABRIC handle, falling back to fd handle exchange\n";
    driver_api->cuMemRelease_(handle);
    return false;
  }

  // 3. check import
  CUmemGenericAllocationHandle import_handle = 0;
  status = driver_api->cuMemImportFromShareableHandle_(
      &import_handle, &sharedHandle, CU_MEM_HANDLE_TYPE_FABRIC);
  if (status != CUDA_SUCCESS) {
    LOG(INFO)
        << "status " << status
        << " Could not import FABRIC handle, falling back to fd handle exchange\n";
    driver_api->cuMemRelease_(handle);
    return false;
  }

  driver_api->cuMemRelease_(import_handle);
  driver_api->cuMemRelease_(handle);
  LOG(INFO) << "using fabric to exchange memory handles\n";
  return true;
}

#endif
} // namespace

bool get_fabric_access(c10::DeviceIndex dev) {
#if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12040 && \
    defined(PYTORCH_C10_DRIVER_API_SUPPORTED)
  TORCH_CHECK(
      num_devices_ >= 0,
      "p2p access cache not initialized. "
      "Ensure c10::cuda::detail::init_p2p_access_cache() is called first.");
  TORCH_CHECK(
      dev >= 0 && dev < num_devices_,
      static_cast<int>(dev),
      " is not a valid device");

  auto& cache = fabricAccessEnabled_[dev];
  if (cache != -1) {
    return cache;
  }

  auto nvml_device = get_nvml_device(dev);
  if (nvml_device != nullptr) {
    nvmlGpuFabricInfoV_t fabricInfo;
    fabricInfo.state = NVML_GPU_FABRIC_STATE_NOT_SUPPORTED;
    fabricInfo.version = nvmlGpuFabricInfo_v2;
    if (DriverAPI::get()->nvmlDeviceGetGpuFabricInfoV_ == nullptr) {
      cache = 0;
      fabricCliqueId_[dev] = kCliqueIdUnsupported;
      return false;
    }
    TORCH_CHECK(
        NVML_SUCCESS ==
        DriverAPI::get()->nvmlDeviceGetGpuFabricInfoV_(
            nvml_device, &fabricInfo));
    auto state = fabricInfo.state != NVML_GPU_FABRIC_STATE_NOT_SUPPORTED;
    if (state) {
      fabricCliqueId_[dev] = static_cast<int>(fabricInfo.cliqueId);
      // now perform the full cycle of allocating - exporting - importing memory
      state = isFabricSupported();
    } else {
      fabricCliqueId_[dev] = kCliqueIdUnsupported;
    }
    cache = state ? 1 : 0;
    return cache;
  } else {
    cache = 0;
    fabricCliqueId_[dev] = kCliqueIdUnsupported;
    return false;
  }
#else
  (void)dev; // Suppress unused parameter warning
  return false;
#endif
}

int get_fabric_clique_id(c10::DeviceIndex dev) {
#if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12040 && \
    defined(PYTORCH_C10_DRIVER_API_SUPPORTED)
  // Ensure cache is populated via get_fabric_access (which does the NVML query
  // and stashes clique_id as a side effect).
  get_fabric_access(dev);
  return fabricCliqueId_[dev];
#else
  (void)dev;
  return kCliqueIdUnsupported;
#endif
}

std::string get_nvml_fabric_info([[maybe_unused]] c10::DeviceIndex dev) {
#if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12040 && \
    defined(PYTORCH_C10_DRIVER_API_SUPPORTED)
  if (DriverAPI::get()->nvmlDeviceGetGpuFabricInfoV_ == nullptr) {
    return "fabric info unsupported (nvmlDeviceGetGpuFabricInfoV not available)";
  }

  auto nvml_device = get_nvml_device(dev);
  if (nvml_device == nullptr) {
    return "fabric info unknown (failed to get NVML device handle)";
  }

  nvmlGpuFabricInfoV_t info{};
#ifdef nvmlGpuFabricInfo_v3
  bool has_health_summary = false;
  info.version = nvmlGpuFabricInfo_v3;
  if (DriverAPI::get()->nvmlDeviceGetGpuFabricInfoV_(nvml_device, &info) ==
      NVML_SUCCESS) {
    has_health_summary = true;
  } else
#endif
  {
    info = {};
    info.version = nvmlGpuFabricInfo_v2;
    if (DriverAPI::get()->nvmlDeviceGetGpuFabricInfoV_(nvml_device, &info) !=
        NVML_SUCCESS) {
      return "fabric info unknown (nvmlDeviceGetGpuFabricInfoV failed)";
    }
  }

  char uuid_hex[33];
  for (int i = 0; i < 16; ++i) {
    snprintf(uuid_hex + i * 2, 3, "%02x", info.clusterUuid[i]);
  }

  const char* state_str = "unknown";
  switch (info.state) {
    case NVML_GPU_FABRIC_STATE_NOT_SUPPORTED:
      state_str = "not_supported";
      break;
    case NVML_GPU_FABRIC_STATE_NOT_STARTED:
      state_str = "not_started";
      break;
    case NVML_GPU_FABRIC_STATE_IN_PROGRESS:
      state_str = "in_progress";
      break;
    case NVML_GPU_FABRIC_STATE_COMPLETED:
      state_str = "completed";
      break;
  }

  std::ostringstream oss;
  oss << "clique_id=" << info.cliqueId << ", cluster_uuid=" << uuid_hex
      << ", state=" << state_str << ", status=" << info.status
      << ", health_mask=0x" << std::hex << std::setfill('0') << std::setw(8)
      << info.healthMask;
#ifdef nvmlGpuFabricInfo_v3
  if (has_health_summary) {
    oss << ", health_summary=" << std::dec
        << static_cast<int>(info.healthSummary);
  }
#endif
  return oss.str();
#else
  return "fabric info unsupported (requires CUDA >= 12.4)";
#endif
}

} // namespace c10::cuda
