/** 
 * Copyright 2021-2023, XGBoost Contributors 
 */

#include <cstdint> // for int64_t
#include "../common/common.h"
#include "../common/device_helpers.cuh" // for DefaultStream, CUDAEvent
#include "array_interface.h"
#include "xgboost/logging.h"

namespace xgboost {

void ArrayInterfaceHandler::SyncCudaStream(std::int64_t stream) {
#if defined(XGBOOST_USE_HIP)
  if (stream == 1 || stream == 2) {
    // Default streams, no synchronization needed
    return;
  }

  hipEvent_t event;
  hipError_t err = hipEventCreate(&event);
  if (err != hipSuccess) {
    LOG(WARNING) << "Failed to create HIP event: " << hipGetErrorString(err);
    return;
  }

  err = hipEventRecord(event, reinterpret_cast<hipStream_t>(stream));
  if (err != hipSuccess) {
    LOG(WARNING) << "Failed to record HIP event: " << hipGetErrorString(err);
    hipEventDestroy(event);
    return;
  }

  err = hipStreamWaitEvent(hipStreamPerThread, event, 0);
  if (err != hipSuccess) {
    LOG(WARNING) << "Failed to wait for HIP event: " << hipGetErrorString(err);
  }

  hipEventDestroy(event);
#else
  switch (stream) {
    case 0:
      /**
       * disallowed by the *`__cuda_array_interface__`*. Quote:
       *
       * This is disallowed as it would be ambiguous between None and the default
       * stream, and also between the legacy and per-thread default streams. Any use
       * case where 0 might be given should either use None, 1, or 2 instead for
       * clarity.
       */
      LOG(FATAL) << "Invalid stream ID in array interface: " << stream;
    case 1: // default legacy stream
      break;
    case 2: // default per-thread stream
    default: {
      dh::CUDAEvent e;
      e.Record(dh::CUDAStreamView{reinterpret_cast<cudaStream_t>(stream)});
      dh::DefaultStream().Wait(e);
    }
  }
#endif
}

bool ArrayInterfaceHandler::IsCudaPtr(void const* ptr) {
  if (!ptr) {
    return false;
  }

  // Check if the pointer is within the process's address space
  uintptr_t ptr_value = reinterpret_cast<uintptr_t>(ptr);
  uintptr_t process_max_addr = (uintptr_t)-1;

#if defined(XGBOOST_USE_HIP)
  hipPointerAttribute_t attr;
  auto err = hipPointerGetAttributes(&attr, ptr);
  

  if (err == hipErrorInvalidValue) {
    return false;
  } else if (err == hipSuccess) {
    switch (attr.type) {
      case hipMemoryTypeUnregistered:
        return false;
      case hipMemoryTypeHost:
        return false;
      case hipMemoryTypeDevice:
        return true;
      case hipMemoryTypeManaged:
        return true;
      default:
        return false;
    }
  } else {
    return false;
  }
#elif defined(XGBOOST_USE_CUDA)
  cudaPointerAttributes attr;
  auto err = cudaPointerGetAttributes(&attr, ptr);
  // reset error
  CHECK_EQ(err, cudaGetLastError());
  if (err == cudaErrorInvalidValue) {
    // CUDA < 11
    return false;
  } else if (err == cudaSuccess) {
    // CUDA >= 11
    switch (attr.type) {
      case cudaMemoryTypeUnregistered:
      case cudaMemoryTypeHost:
        return false;
      default:
        return true;
    }
    return true;
  } else {
    // other errors, `cudaErrorNoDevice`, `cudaErrorInsufficientDriver` etc.
    return false;
  }
#else
  return false;
#endif
}

}  // namespace xgboost