/*
 * SPDX-FileCopyrightText: Copyright (c) 2021-2024 NVIDIA CORPORATION &
 * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
 *
 * 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.
 */
#ifndef NVIDIA_GXF_CUDA_CUDA_COMMON_HPP_
#define NVIDIA_GXF_CUDA_CUDA_COMMON_HPP_

#include <string>
#include <cuda_runtime.h>

#include "common/assert.hpp"
#include "common/logger.hpp"

#ifdef _MSC_VER

#include "gxf/core/gxf.h"

#define CHECK_CUDA_ERROR_2(cu_result, fmt)                                        \
  do {                                                                            \
    cudaError_t err = (cu_result);                                                \
    if (err != cudaSuccess) {                                                     \
      GXF_LOG_ERROR(fmt ", cuda_error: %s, error_str: %s", cudaGetErrorName(err), \
                    cudaGetErrorString(err));                                     \
      return Unexpected{GXF_FAILURE};                                             \
    }                                                                             \
  } while (0)

#define CHECK_CUDA_ERROR_3(cu_result, fmt, ...)                                                \
  do {                                                                                         \
    cudaError_t err = (cu_result);                                                             \
    if (err != cudaSuccess) {                                                                  \
      GXF_LOG_ERROR(fmt ", cuda_error: %s, error_str: %s", __VA_ARGS__, cudaGetErrorName(err), \
                    cudaGetErrorString(err));                                                  \
      return Unexpected{GXF_FAILURE};                                                          \
    }                                                                                          \
  } while (0)

// Dispatcher macro
#define GET_MACRO(_1, _2, _3, NAME, ...) NAME

#define CHECK_CUDA_ERROR(...) \
  GET_MACRO(__VA_ARGS__, CHECK_CUDA_ERROR_3, CHECK_CUDA_ERROR_2)(__VA_ARGS__)

inline auto check_cuda_error_result_fn(cudaError_t err, const char* fmt) {
  if (err != cudaSuccess) {
    GXF_LOG_ERROR("%s, cuda_error: %s, error_str: %s", fmt, cudaGetErrorName(err),
                  cudaGetErrorString(err));
    return GXF_FAILURE;
  }
  return GXF_SUCCESS;
}

template <typename... Args>
inline auto check_cuda_error_result_fn(cudaError_t err, const char* fmt, Args&&... args) {
  if (err != cudaSuccess) {
    std::string formatted;
    int len = std::snprintf(nullptr, 0, fmt, std::forward<Args>(args)...);
    if (len > 0) {
      formatted.resize(len + 1);
      std::snprintf(&formatted[0], formatted.size(), fmt, std::forward<Args>(args)...);
      formatted.pop_back();  // remove trailing null
    }
    GXF_LOG_ERROR("%s, cuda_error: %s, error_str: %s", formatted.c_str(), cudaGetErrorName(err),
                  cudaGetErrorString(err));
    return GXF_FAILURE;
  }
  return GXF_SUCCESS;
}

#define CHECK_CUDA_ERROR_RESULT(cu_result, ...) check_cuda_error_result_fn((cu_result), __VA_ARGS__)

template <typename... Args>
void continue_cuda_error(cudaError_t err, const char* fmt, Args&&... args) {
  if (err != cudaSuccess) {
    std::string full_fmt = std::string(fmt) + ", cuda_error: {}, error_str: {}";
    GXF_LOG_ERROR(full_fmt.c_str(), std::forward<Args>(args)..., cudaGetErrorName(err),
                  cudaGetErrorString(err));
  }
}

#define CONTINUE_CUDA_ERROR(cu_result, ...) continue_cuda_error((cu_result), __VA_ARGS__)

#else
#define CHECK_CUDA_ERROR(cu_result, fmt, ...)                                  \
    do {                                                                       \
        cudaError_t err = (cu_result);                                          \
        if (err != cudaSuccess) {                                               \
            GXF_LOG_ERROR(fmt ", cuda_error: %s, error_str: %s", ##__VA_ARGS__, \
                cudaGetErrorName(err), cudaGetErrorString(err));                \
            return Unexpected{GXF_FAILURE};                                     \
        }                                                                       \
    } while (0)

#define CHECK_CUDA_ERROR_RESULT(cu_result, fmt, ...)                      \
  do {                                                                    \
    cudaError_t err = (cu_result);                                        \
    if (err != cudaSuccess) {                                             \
      GXF_LOG_ERROR(fmt ", cuda_error: %s, error_str: %s", ##__VA_ARGS__, \
                    cudaGetErrorName(err), cudaGetErrorString(err));      \
      return GXF_FAILURE;                                                 \
    }                                                                     \
  } while (0)

#define CONTINUE_CUDA_ERROR(cu_result, fmt, ...)                               \
    do {                                                                       \
        cudaError_t err = (cu_result);                                          \
        if (err != cudaSuccess) {                                               \
            GXF_LOG_ERROR(fmt ", cuda_error: %s, error_str: %s", ##__VA_ARGS__, \
                cudaGetErrorName(err), cudaGetErrorString(err));                \
        }                                                                       \
    } while (0)
#endif // _MSC_VER

#endif  // NVIDIA_GXF_CUDA_CUDA_COMMON_HPP_
