/*
 * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved.
 *
 * This source code and/or documentation ("Licensed Deliverables") are
 * subject to NVIDIA intellectual property rights under U.S. and
 * international Copyright laws.
 */

#pragma once

#include <cugraph-ops/cuda/stream.hpp>
#include <cugraph-ops/utils/allocator.hpp>
#include <cugraph-ops/utils/logger.hpp>

#include <utils/cudart.hpp>

#include <cinttypes>
#include <type_traits>

namespace cugraph::ops {

/**
 * @brief Type-safe RAII abstraction for all pointers used on device
 *
 * @note Currently, this class assumes the total ownership of the lifetime of
 *       the underlying pointer.
 */
template <typename DataT, cudaMemoryType PTYPE>
class __attribute__((aligned(16))) test_buffer {
 public:
  static constexpr cudaMemoryType TYPE = PTYPE;
  using data_t                         = DataT;
  test_buffer()                        = default;
  /**
   * @brief Allocate a buffer of the given length
   *
   * @param[in] length length of the buffer in number of elements, limited by INT_MAX
   * @param[in] reset  if set, buffer is set to all zeros
   * @param[in] stream stream to which the buffer's lifetime is associated, NULL stream by default
   */
  template <typename IntT, std::enable_if_t<std::is_integral_v<IntT>, bool> = true>
  explicit test_buffer(IntT length, bool reset = false, cudaStream_t stream = nullptr)
    : stream_(stream)
  {
    if (static_cast<int>(length) < 0) length = 0;
    reserve(static_cast<size_t>(length), reset);
  }

  template <typename DataBT, typename IntT, std::enable_if_t<std::is_integral_v<IntT>, bool> = true>
  test_buffer(const DataBT* other, IntT length, cudaStream_t stream = nullptr) : stream_(stream)
  {
    if (static_cast<int>(length) < 0) length = 0;
    reserve(static_cast<size_t>(length));
    if (length <= 0) return;
    copy_from(other, length);
  };

  template <typename DataBT, cudaMemoryType PTYPE2>
  explicit test_buffer(const test_buffer<DataBT, PTYPE2>& other) : stream_(other.cuda_stream())
  {
    reserve(other.length());
    if (other.length() <= 0) return;
    copy_from(other);
  };

  /**
   * @brief Destroys the buffer
   */
  ~test_buffer() noexcept { dealloc(); }

  void dealloc() noexcept
  {
    try {
      if (data_ != nullptr) {
        if (PTYPE == cudaMemoryTypeDevice)
          utils::device_allocator().dealloc(data_, size(), cuda_stream());
        else if (PTYPE == cudaMemoryTypeHost)
          utils::host_allocator().dealloc(data_, size(), cuda_stream());
        else if (PTYPE == cudaMemoryTypeManaged)
          utils::managed_allocator().dealloc(data_, size(), cuda_stream());
        data_ = nullptr;
        size_ = 0;
      }
    } catch (const std::exception& e) {
      CUGRAPH_OPS_ERROR("Got exception during de-allocation: %s", e.what());
    }
  }

  void reserve(size_t new_length, bool reset = false) noexcept
  {
    auto new_size = new_length * sizeof(DataT);
    if (size() > 0) dealloc();
    if (new_size <= 0) return;
    size_ = new_size;
    try {
      if (PTYPE == cudaMemoryTypeDevice) {
        data_ = reinterpret_cast<DataT*>(utils::device_allocator().alloc(size(), cuda_stream()));
      } else if (PTYPE == cudaMemoryTypeHost) {
        data_ = reinterpret_cast<DataT*>(utils::host_allocator().alloc(size(), cuda_stream()));
      } else {
        data_ = reinterpret_cast<DataT*>(utils::managed_allocator().alloc(size(), cuda_stream()));
        cudaPointerAttributes attr;
        auto ret = cudaPointerGetAttributes(&attr, data_);
        if (ret == cudaSuccess && attr.hostPointer != nullptr) {
          data_ = reinterpret_cast<DataT*>(attr.hostPointer);
        } else {
          CUGRAPH_OPS_ERROR("Unable to get host pointer from managed allocation");
        }
      }
    } catch (const std::exception& e) {
      CUGRAPH_OPS_ERROR("Got exception during allocation: %s", e.what());
      return;
    }
    if (reset) {
      utils::memset_async_no_throw(data_, new_length, cuda_stream());
    } else {
      static constexpr uint8_t RND_BYTE = 101;
      utils::memset_async_no_throw(data_, new_length, cuda_stream(), RND_BYTE);
    }
    // sync to ensure that this buffer is valid
    utils::sync_no_throw(cuda_stream());
  }

  /**
   * @brief Returns the underlying raw pointer
   *
   * @note: not marked explicit to simplify test usage
   * @{
   */
  // NOLINTNEXTLINE(google-explicit-constructor)
  operator DataT*() const noexcept { return data_; }
  [[nodiscard]] DataT* data() noexcept { return data_; }
  [[nodiscard]] const DataT* data() const noexcept { return data_; }
  /** @} */

  /** Returns the size of the buffer (in bytes) */
  [[nodiscard]] size_t size() const noexcept { return size_; }

  /** Returns the size of the buffer (in elements) */
  [[nodiscard]] size_t length() const noexcept { return size_ / sizeof(DataT); }

  /** Checks whether the buffer is empty */
  [[nodiscard]] bool empty() const noexcept { return 0 == size(); }

  /** Returns the underlying cuda stream */
  [[nodiscard]] cudaStream_t cuda_stream() const noexcept { return stream_; }

  /** Access elements on non-device buffers */
  template <cudaMemoryType PT = PTYPE>
  [[nodiscard]] std::enable_if_t<PT != cudaMemoryTypeDevice, DataT&> operator[](int i) noexcept
  {
    return data_[i];
  }
  template <cudaMemoryType PT = PTYPE>
  [[nodiscard]] std::enable_if_t<PT != cudaMemoryTypeDevice, const DataT&> operator[](
    int i) const noexcept
  {
    return data_[i];
  }

  // we delete move/copy assignment because we want to allow copying from
  // different data types and pointer types when possible
  test_buffer(test_buffer&& other)            = delete;
  test_buffer& operator=(test_buffer&& other) = delete;

  test_buffer(const test_buffer& other)            = delete;
  test_buffer& operator=(const test_buffer& other) = delete;

  template <typename DataBT, typename IntT, std::enable_if_t<std::is_integral_v<IntT>, bool> = true>
  void copy_from(const DataBT* other, IntT len) noexcept
  {
    if (len > length()) {
      const auto* err = "Cannot copy buffer of length %" PRId64 " to buffer of length %" PRId64;
      CUGRAPH_OPS_ERROR(err, static_cast<int64_t>(len), static_cast<int64_t>(length()));
      return;
    }
    const DataBT* h_other = other;
    // we are able to rely on this function to detect device pointers, because
    // they have to be alloc'ed with cudaMalloc (or similar) and we only support
    // GPUs which support unified addressing
    cudaPointerAttributes attr;
    auto ret          = cudaPointerGetAttributes(&attr, other);
    bool is_other_dev = ret == cudaSuccess && attr.type == cudaMemoryTypeDevice;
    if (!is_other_dev && attr.hostPointer != nullptr) {
      h_other = reinterpret_cast<const DataBT*>(attr.hostPointer);
    }

    if constexpr (std::is_same_v<DataT, DataBT>) {
      auto kind       = is_other_dev && PTYPE == cudaMemoryTypeDevice ? cudaMemcpyDeviceToDevice
                        : is_other_dev                                ? cudaMemcpyDeviceToHost
                        : PTYPE == cudaMemoryTypeDevice               ? cudaMemcpyHostToDevice
                                                                      : cudaMemcpyHostToHost;
      auto* other_ptr = is_other_dev ? other : h_other;
      utils::copy_async_no_throw(data(), other_ptr, len, kind, cuda_stream());
    } else if constexpr (PTYPE == cudaMemoryTypeDevice) {
      test_buffer<DataBT, cudaMemoryTypeHost> h_int1{};
      test_buffer<DataT, cudaMemoryTypeHost> h_int2{len};

      if (is_other_dev) {
        h_int1.reserve(len);
        utils::copy_async_no_throw<DataBT>(
          h_int1, other, len, cudaMemcpyDeviceToHost, cuda_stream());
        utils::sync_no_throw(cuda_stream());
        h_other = h_int1;
      }
      for (size_t i = 0; i < len; ++i)
        h_int2[i] = static_cast<DataT>(h_other[i]);

      auto kind = PTYPE == cudaMemoryTypeDevice ? cudaMemcpyHostToDevice : cudaMemcpyHostToHost;
      utils::copy_async_no_throw<DataT>(data_, h_int2, len, kind, cuda_stream());
    } else {
      test_buffer<DataBT, cudaMemoryTypeHost> h_int1{len};
      auto kind       = is_other_dev ? cudaMemcpyDeviceToHost : cudaMemcpyHostToHost;
      auto* other_ptr = is_other_dev ? other : h_other;
      utils::copy_async_no_throw<DataBT>(h_int1, other_ptr, len, kind, cuda_stream());
      utils::sync_no_throw(cuda_stream());
      for (size_t i = 0; i < len; ++i)
        data_[i] = DataT{h_int1[i]};
    }
    // sync to ensure that this buffer has valid data
    utils::sync_no_throw(cuda_stream());
  }
  template <typename DataBT, cudaMemoryType PTYPE2>
  void copy_from(const test_buffer<DataBT, PTYPE2>& other) noexcept
  {
    if (other.length() > length()) {
      const auto* err =
        "Cannot copy test_buffer of length %" PRId64 " to buffer of length %" PRId64;
      CUGRAPH_OPS_ERROR(err, static_cast<int64_t>(other.length()), static_cast<int64_t>(length()));
      return;
    }
    // sync to ensure that other buffer has valid data
    if (other.cuda_stream() != cuda_stream()) utils::sync_no_throw(other.cuda_stream());
    copy_from(other.data(), other.length());
  }

 private:
  /** size of the buffer (in B) */
  size_t size_{0};
  /** underlying cuda stream */
  cudaStream_t stream_{nullptr};
  /** pointer to the buffer */
  DataT* data_{nullptr};
};  // class test_buffer

// these aliases should represent structs, rather than simple type aliases
// NOLINTBEGIN(readability-identifier-naming)
template <typename DataT>
using device_buffer = test_buffer<DataT, cudaMemoryTypeDevice>;
template <typename DataT>
using host_buffer = test_buffer<DataT, cudaMemoryTypeHost>;
template <typename DataT>
using managed_buffer = test_buffer<DataT, cudaMemoryTypeManaged>;
// NOLINTEND(readability-identifier-naming)

}  // namespace cugraph::ops
