// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#pragma once

#include <cub/config.cuh>

#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
#  pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
#  pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
#  pragma system_header
#endif // no system header

#include <cub/util_device.cuh>

#include <thrust/system/cuda/detail/core/triple_chevron_launch.h>

CUB_NAMESPACE_BEGIN

namespace detail
{
struct TripleChevronFactory
{
  CUB_RUNTIME_FUNCTION THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron operator()(
    dim3 grid, dim3 block, ::cuda::std::size_t shared_mem, ::cudaStream_t stream, bool dependent_launch = false) const
  {
    return THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(grid, block, shared_mem, stream, dependent_launch);
  }

  template <class T = void>
  CUB_RUNTIME_FUNCTION ::cudaError_t PtxVersion(int& version)
  {
    return cub::PtxVersion<T>(version);
  }

  template <class T = void>
  CUB_RUNTIME_FUNCTION ::cudaError_t PtxArchId(::cuda::arch_id& arch_id) const
  {
    return ptx_arch_id<T>(arch_id);
  }

  _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t MultiProcessorCount(int& sm_count) const
  {
    int device_ordinal;
    ::cudaError_t error = CubDebug(::cudaGetDevice(&device_ordinal));
    if (::cudaSuccess != error)
    {
      return error;
    }

    // Get SM count
    return ::cudaDeviceGetAttribute(&sm_count, ::cudaDevAttrMultiProcessorCount, device_ordinal);
  }

  template <typename Kernel>
  _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t
  MaxSmOccupancy(int& sm_occupancy, Kernel kernel_ptr, int block_size, int dynamic_smem_bytes = 0)
  {
    return ::cudaOccupancyMaxActiveBlocksPerMultiprocessor(&sm_occupancy, kernel_ptr, block_size, dynamic_smem_bytes);
  }

  _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t MaxGridDimX(int& max_grid_dim_x) const
  {
    int device_ordinal;
    ::cudaError_t error = CubDebug(::cudaGetDevice(&device_ordinal));
    if (::cudaSuccess != error)
    {
      return error;
    }

    // Get max grid dimension
    return ::cudaDeviceGetAttribute(&max_grid_dim_x, ::cudaDevAttrMaxGridDimX, device_ordinal);
  }

  _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t
  MemsetAsync(void* dst, unsigned char value, size_t num_bytes, ::cudaStream_t stream) const
  {
    return ::cudaMemsetAsync(dst, value, num_bytes, stream);
  }

  _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t
  MemcpyAsync(void* dst, const void* src, size_t num_bytes, ::cudaMemcpyKind kind, ::cudaStream_t stream) const
  {
    return ::cudaMemcpyAsync(dst, src, num_bytes, kind, stream);
  }

  // TODO(bgruber): this is very similar to thrust::cuda_cub::core::get_max_shared_memory_per_block. We should unify
  // this.
  _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION cudaError_t MaxSharedMemory(int& max_shared_memory) const
  {
    int device = 0;
    auto error = CubDebug(cudaGetDevice(&device));
    if (error != cudaSuccess)
    {
      return error;
    }

    return cudaDeviceGetAttribute(&max_shared_memory, cudaDevAttrMaxSharedMemoryPerBlock, device);
  }

  template <typename Kernel>
  _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t
  max_dynamic_smem_size_for(int& max_dynamic_smem_size, [[maybe_unused]] Kernel kernel_ptr)
  {
    NV_IF_ELSE_TARGET(NV_IS_HOST, //
                      ({ return MaxPotentialDynamicSmemBytes(max_dynamic_smem_size, kernel_ptr); }),
                      ({
                        ::cudaFuncAttributes func_attrs{};
                        if (const auto error = CubDebug(::cudaFuncGetAttributes(&func_attrs, kernel_ptr)))
                        {
                          return error;
                        }
                        max_dynamic_smem_size = func_attrs.maxDynamicSharedSizeBytes;
                        return cudaSuccess;
                      }))
  }

  template <typename Kernel>
  _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t set_max_dynamic_smem_size_for(Kernel kernel_ptr, int smem_size)
  {
    return CubDebug(::cudaFuncSetAttribute(kernel_ptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
  }
};
} // namespace detail

CUB_NAMESPACE_END
