/*
 * Copyright (c) 2020-2022, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

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

#include <cuda_runtime.h>

#include <cstring>
#include <functional>

namespace cugraph::ops::gnn::utils {

/** functor for allocating memory */
using alloc_func_t = std::function<void*(size_t, const cuda::stream&)>;
/** functor for deallocating memory */
using dealloc_func_t = std::function<void(void*, size_t, const cuda::stream&)>;

/**
 * @brief A simple wrapper for allocating memory (either host or device)
 */
struct allocator {
  /** functor used to allocate objects */
  alloc_func_t alloc;
  /** functor used to deallocate objects */
  dealloc_func_t dealloc;
};  // struct allocator

/**
 * @brief Returns reference to tbe stored device allocator object
 */
allocator& device_allocator();

/**
 * @brief Sets the input object as the device allocator for the subsequent calls
 *
 * @param[in] obj allocator to be used as the default for all subsequent calls
 *
 * @return the previous allocator
 *
 * @note Setting a new allocator without clearing the memory allocated via the
 *       old one will cause memory leaks!
 */
allocator set_device_allocator(const allocator& obj);

/**
 * @brief Returns reference to tbe stored host (aka pinned memory) allocator
 *        object
 */
allocator& host_allocator();

/**
 * @brief Sets the input object as the host allocator for the subsequent calls
 *
 * @param[in] obj allocator to be used as the default for all subsequent calls
 *
 * @return the previous allocator
 *
 * @note Setting a new allocator without clearing the memory allocated via the
 *       old one will cause memory leaks!
 */
allocator set_host_allocator(const allocator& obj);

/**
 * @brief Returns reference to tbe stored managed memory allocator object
 */
allocator& managed_allocator();

/**
 * @brief Sets the input object as the managed memory allocator for the
 *        subsequent calls
 *
 * @param[in] obj allocator to be used as the default for all subsequent calls
 *
 * @return the previous allocator
 *
 * @note Setting a new allocator without clearing the memory allocated via the
 *       old one will cause memory leaks!
 */
allocator set_managed_allocator(const allocator& obj);

/**
 * @brief Returns reference to tbe stored allocator
 *        object which is specified by the desired memory type
 *
 * @param[in] type of memory
 */
allocator generic_allocator(cudaMemoryType type);

/**
 * @brief allocate a (pinned) host buffer
 *
 * @tparam DataT data type
 *
 * @param[out] ptr   output array
 * @param[in]  len   length of the array
 * @param[in]  s     cuda stream
 * @param[in]  reset reset the array to zeros
 */
template <typename DataT>
void host_alloc(DataT*& ptr, size_t len, cudaStream_t s, bool reset = false)
{
  ptr = reinterpret_cast<DataT*>(host_allocator().alloc(sizeof(DataT) * len, s));
  if (reset) { memset(ptr, 0, sizeof(DataT) * len); }
}

/**
 * @brief deallocate a (pinned) host buffer
 *
 * @tparam DataT data type
 *
 * @param[in] ptr array to be deallocated
 * @param[in]  s  cuda stream
 */
template <typename DataT>
void host_free(DataT* ptr, cudaStream_t s)
{
  host_allocator().dealloc(ptr, 0, s);
}

/**
 * @brief allocate a buffer on gpu memory
 *
 * @tparam DataT data type
 *
 * @param[out] ptr   output array
 * @param[in]  len   length of the array
 * @param[in]  s     cuda stream
 * @param[in]  reset reset the array to zeros
 */
template <typename DataT>
void device_alloc(DataT*& ptr, size_t len, cudaStream_t s, bool reset = false)
{
  ptr = reinterpret_cast<DataT*>(device_allocator().alloc(sizeof(DataT) * len, s));
  if (reset) { RAFT_CUDA_TRY(cudaMemsetAsync(ptr, 0, sizeof(DataT) * len, s)); }
}

/**
 * @brief deallocate a gpu buffer
 *
 * @tparam DataT data type
 *
 * @param[in] ptr array to be deallocated
 * @param[in]  s  cuda stream
 */
template <typename DataT>
void device_free(DataT* ptr, cudaStream_t s)
{
  device_allocator().dealloc(ptr, 0, s);
}

/**
 * @brief allocate a managed buffer
 *
 * @tparam DataT data type
 *
 * @param[out] ptr   output array
 * @param[in]  len   length of the array
 * @param[in]  s     cuda stream
 * @param[in]  reset reset the array to zeros
 */
template <typename DataT>
void managed_alloc(DataT*& ptr, size_t len, cudaStream_t s, bool reset = false)
{
  ptr = reinterpret_cast<DataT*>(managed_allocator().alloc(sizeof(DataT) * len, s));
  if (reset) { memset(ptr, 0, sizeof(DataT) * len); }  // reset on CPU
}

/**
 * @brief deallocate a managed buffer
 *
 * @tparam DataT data type
 *
 * @param[in] ptr array to be deallocated
 * @param[in]  s  cuda stream
 */
template <typename DataT>
void managed_free(DataT* ptr, cudaStream_t s)
{
  managed_allocator().dealloc(ptr, 0, s);
}

}  // namespace cugraph::ops::gnn::utils
