// SPDX-FileComment: This file is part of TNL - Template Numerical Library (https://tnl-project.org/)
// SPDX-License-Identifier: MIT

#pragma once

#include <cstdint>
#include <iostream>
#include <limits>       // std::numeric_limits
#include <memory>       // std::unique_ptr
#include <type_traits>  // std::remove_cv_t

#include "Types.h"
#include "Functions.h"
#include <TNL/DiscreteMath.h>
#include <TNL/Math.h>

namespace TNL::Backend {

constexpr std::size_t
getMaxGridXSize()
{
#if defined( __HIPCC__ )
   // AMD GPUs support up to 2^32 - 1 blocks in the x-dimension
   // https://rocm.docs.amd.com/projects/HIP/en/latest/reference/hardware_features.html
   return 4294967295ULL;
#else
   return 2147483647;
#endif
}

constexpr std::size_t
getMaxGridYSize()
{
   // Both NVIDIA and AMD GPUs have the same limit of 65535 blocks
   // https://rocm.docs.amd.com/projects/HIP/en/latest/reference/hardware_features.html
   return 65535;
}

constexpr std::size_t
getMaxGridZSize()
{
   // Both NVIDIA and AMD GPUs have the same limit of 65535 blocks
   // https://rocm.docs.amd.com/projects/HIP/en/latest/reference/hardware_features.html
   return 65535;
}

constexpr std::size_t
getMaxThreadsPerGrid()
{
#if defined( __HIPCC__ )
   // AMD GPUs limit total threads per grid launch to 2^32 - 1
   // https://rocm.docs.amd.com/projects/HIP/en/latest/reference/hardware_features.html
   return 4294967295ULL;
#else
   // NVIDIA GPUs have no explicit threads-per-grid limit beyond the grid/block
   // dimension limits
   return std::numeric_limits< std::size_t >::max();
#endif
}

constexpr int
getMaxThreadsPerBlock()
{
   // Both NVIDIA and AMD GPUs support up to 1024 threads per block
   return 1024;
}

constexpr int
getMaxBlockXSize()
{
   return 1024;
}

constexpr int
getMaxBlockYSize()
{
   return 1024;
}

constexpr int
getMaxBlockZSize()
{
#if defined( __CUDACC__ ) || defined( __HIP_PLATFORM_NVCC__ )
   return 64;
#else
   return 1024;
#endif
}

/**
 * The warpSize variable is of type int and contains the warp size (in threads)
 * for the target device. This should be used only from device code in order to
 * develop portable wave-aware code.
 *
 * Note that NVIDIA devices return 32; AMD devices return 64 for gfx8/gfx9 and
 * 32 for gfx10 and above.
 *
 * Warning: the returned value may be inconsistent when used from a host-code
 * context on HIP, because `__GFX8__`/`__GFX9__` are defined only during device
 * compilation. On the host side, the function falls through to return 32.
 * Use `getMaxWarpSize()` for compile-time host-side guards, or the runtime
 * overload `getWarpSize(int deviceId)` for host-side dispatch decisions.
 * See https://clang.llvm.org/docs/HIPSupport.html#predefined-macros for details.
 *
 * https://rocm.docs.amd.com/projects/HIP/en/latest/reference/kernel_language.html#warpsize
 */
constexpr int
getWarpSize()
{
#if defined( __CUDACC__ ) || defined( __HIP_PLATFORM_NVCC__ )
   return 32;
// We use architecture macros instead of the deprecated `__AMDGCN_WAVEFRONT_SIZE__`
// which was removed in ROCm 7.0+.
#elif defined( __GFX8__ ) || defined( __GFX9__ )
   // gfx8/gfx9 (GCN) use wave64, all others use wave32
   return 64;
#else
   return 32;
#endif
}

/**
 * \brief Returns the maximum possible warp/wavefront size for the current build.
 *
 * Unlike getWarpSize(), this is consistent between host and device compilation:
 * - CUDA: always 32
 * - HIP: always 64 (the maximum across all AMD architectures)
 * - Host-only: 32
 *
 * Use this for compile-time guards that must be consistent on the host side,
 * e.g. deciding whether to instantiate TPS=64 kernel variants.
 */
constexpr int
getMaxWarpSize()
{
#if defined( __CUDACC__ ) || defined( __HIP_PLATFORM_NVCC__ )
   return 32;
#elif defined( __HIP__ )
   return 64;
#else
   return 32;
#endif
}

/**
 * \brief Returns the minimum possible warp/wavefront size across all platforms.
 *
 * Always returns 32. Use this for compile-time checks that must allow configs
 * valid on any architecture, e.g. column-major SlicedEllpack launch configs
 * where `TPS * SliceSize >= getMinWarpSize()` ensures the group covers at
 * least one warp on the smallest-warp-size device.
 */
constexpr int
getMinWarpSize()
{
   return 32;
}

/**
 * \brief Runtime query for the warp/wavefront size of a specific device.
 *
 * Use this for host-side launch configuration decisions (TPS selection, etc.).
 * For CUDA, always returns 32. For HIP, queries the runtime API.
 *
 * \param deviceId The device ID to query (default: current device).
 */
[[nodiscard]] inline int
getWarpSize( int deviceId )
{
#if defined( __CUDACC__ ) || defined( __HIP_PLATFORM_NVCC__ )
   (void) deviceId;
   return 32;
#elif defined( __HIP__ )
   int warpSize = 0;
   TNL_BACKEND_SAFE_CALL( hipDeviceGetAttribute( &warpSize, hipDeviceAttributeWarpSize, deviceId ) );
   return warpSize;
#else
   (void) deviceId;
   return 32;
#endif
}

/**
 * \brief Returns the full mask for warp shuffle operations.
 *
 * HIP shfl intrinsics require a 64-bit mask regardless of the wavefront size
 * (unused upper bits are zero for wave32). CUDA uses a 32-bit mask.
 *
 * Warning: this function relies on `__GFX8__`/`__GFX9__` arch macros which are
 * defined only during device compilation. On the host side, it falls through
 * to the wavefront=32 mask. Do not call from host code — use only in
 * `__device__` functions or inside `#if defined(__CUDACC__) || defined(__HIP__)`
 * guards.
 */
[[nodiscard]] constexpr auto
getWarpFullMask()
{
#if defined( __CUDACC__ ) || defined( __HIP_PLATFORM_NVCC__ )
   return std::uint32_t{ 0xffffffff };
#elif defined( __GFX8__ ) || defined( __GFX9__ )
   // gfx8/gfx9 (GCN) use wave64, all others use wave32
   return std::uint64_t{ 0xffffffffffffffffULL };
#else
   return std::uint64_t{ 0xffffffffULL };
#endif
}

// When we transfer data between the GPU and the CPU we use 1 MiB buffer. This
// size should ensure good performance.
// We use the same buffer size even for retyping data during IO operations.
constexpr std::size_t
getTransferBufferSize()
{
   return 1 << 20;
}

template< typename Element, typename FillBuffer, typename PushBuffer >
void
bufferedTransfer( std::size_t size, FillBuffer& fill, PushBuffer& push )
{
   const std::size_t buffer_size = std::min( Backend::getTransferBufferSize() / sizeof( Element ), size );
   std::unique_ptr< Element[] > host_buffer{ new Element[ buffer_size ] };

   std::size_t compared = 0;
   while( compared < size ) {
      const std::size_t transfer = std::min( size - compared, buffer_size );
      fill( compared, host_buffer.get(), transfer );
      bool next_iter = true;
      push( compared, host_buffer.get(), transfer, next_iter );
      if( ! next_iter )
         return;
      compared += transfer;
   }
}

/**
 * This function creates a buffer on the host, fills it with data transferred
 * from \e source, which is a pointer to device memory, and the \e push handler
 * processes the data in the buffer.
 */
template< typename Element, typename PushBuffer >
void
bufferedTransferToHost( const Element* source, std::size_t size, PushBuffer& push )
{
   using BufferType = std::remove_cv_t< Element >;
   auto fill = [ source ]( std::size_t offset, BufferType* buffer, std::size_t buffer_size )
   {
      Backend::memcpy(
         static_cast< void* >( buffer ),
         static_cast< const void* >( source + offset ),
         buffer_size * sizeof( Element ),
         Backend::MemcpyDeviceToHost );
   };
   bufferedTransfer< BufferType >( size, fill, push );
}

/**
 * This function creates a buffer on the host, the \e fill handler fills it
 * with data and this function transfers data from the buffer to the
 * \e destination, which is a pointer to device memory.
 */
template< typename Element, typename FillBuffer >
void
bufferedTransferToDevice( Element* destination, std::size_t size, FillBuffer& fill )
{
   using BufferType = std::remove_cv_t< Element >;
   auto push = [ destination ]( std::size_t offset, const BufferType* buffer, std::size_t buffer_size, bool& next_iter )
   {
      Backend::memcpy(
         static_cast< void* >( destination + offset ),
         static_cast< const void* >( buffer ),
         buffer_size * sizeof( Element ),
         Backend::MemcpyHostToDevice );
   };
   bufferedTransfer< BufferType >( size, fill, push );
}

#if defined( __CUDACC__ ) || defined( __HIP__ )
__device__
inline std::size_t
getGlobalBlockIdx_x( const dim3& gridIdx )
{
   return static_cast< std::size_t >( gridIdx.x ) * getMaxGridXSize() + blockIdx.x;
}

__device__
inline std::size_t
getGlobalBlockIdx_y( const dim3& gridIdx )
{
   return static_cast< std::size_t >( gridIdx.y ) * getMaxGridYSize() + blockIdx.y;
}

__device__
inline std::size_t
getGlobalBlockIdx_z( const dim3& gridIdx )
{
   return static_cast< std::size_t >( gridIdx.z ) * getMaxGridZSize() + blockIdx.z;
}

__device__
inline std::size_t
getGlobalThreadIdx_x( const dim3& gridIdx )
{
   return getGlobalBlockIdx_x( gridIdx ) * blockDim.x + threadIdx.x;
}

__device__
inline std::size_t
getGlobalThreadIdx_y( const dim3& gridIdx )
{
   return getGlobalBlockIdx_y( gridIdx ) * blockDim.y + threadIdx.y;
}

__device__
inline std::size_t
getGlobalThreadIdx_z( const dim3& gridIdx )
{
   return getGlobalBlockIdx_z( gridIdx ) * blockDim.z + threadIdx.z;
}
#endif

template< typename Index >
Index
getNumberOfBlocks( const Index threads, const unsigned int blockSize )
{
   // Index may be actually smaller type than unsigned int
   using CommonType = std::common_type_t< Index, unsigned int >;
   return roundUpDivision( static_cast< CommonType >( threads ), static_cast< CommonType >( blockSize ) );
}

template< typename Index >
Index
getNumberOfGrids( const Index blocks, const unsigned int gridSize )
{
   // Index may be actually smaller type than unsigned int
   using CommonType = std::common_type_t< Index, unsigned int >;
   return roundUpDivision( static_cast< CommonType >( blocks ), static_cast< CommonType >( gridSize ) );
}

inline void
setupThreads(
   const dim3& blockSize,
   dim3& blocksCount,
   dim3& gridsCount,
   long long int xThreads,
   long long int yThreads = 0,
   long long int zThreads = 0 )
{
   blocksCount.x = max( 1, xThreads / blockSize.x + static_cast< long long int >( xThreads % blockSize.x != 0 ) );
   blocksCount.y = max( 1, yThreads / blockSize.y + static_cast< long long int >( yThreads % blockSize.y != 0 ) );
   blocksCount.z = max( 1, zThreads / blockSize.z + static_cast< long long int >( zThreads % blockSize.z != 0 ) );

   gridsCount.x =
      blocksCount.x / getMaxGridXSize() + static_cast< unsigned long int >( blocksCount.x % getMaxGridXSize() != 0 );
   gridsCount.y =
      blocksCount.y / getMaxGridYSize() + static_cast< unsigned long int >( blocksCount.y % getMaxGridYSize() != 0 );
   gridsCount.z =
      blocksCount.z / getMaxGridZSize() + static_cast< unsigned long int >( blocksCount.z % getMaxGridZSize() != 0 );
}

inline void
setupGrid( const dim3& blocksCount, const dim3& gridsCount, const dim3& gridIdx, dim3& gridSize )
{
   if( gridIdx.x < gridsCount.x - 1 )
      gridSize.x = getMaxGridXSize();
   else
      gridSize.x = blocksCount.x % getMaxGridXSize();

   if( gridIdx.y < gridsCount.y - 1 )
      gridSize.y = getMaxGridYSize();
   else
      gridSize.y = blocksCount.y % getMaxGridYSize();

   if( gridIdx.z < gridsCount.z - 1 )
      gridSize.z = getMaxGridZSize();
   else
      gridSize.z = blocksCount.z % getMaxGridZSize();
}

inline std::ostream&
operator<<( std::ostream& str, const dim3& d )
{
   str << "( " << d.x << ", " << d.y << ", " << d.z << " )";
   return str;
}

inline void
printThreadsSetup(
   const dim3& blockSize,
   const dim3& blocksCount,
   const dim3& gridSize,
   const dim3& gridsCount,
   std::ostream& str = std::cout )
{
   str << "Block size: " << blockSize << "\n Blocks count: " << blocksCount << "\n Grid size: " << gridSize
       << "\n Grids count: " << gridsCount << "\n";
}

}  // namespace TNL::Backend
