//===----------------------------------------------------------------------===//
//
// Part of libcu++, the C++ Standard Library for your entire system,
// under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//

#ifndef __CUDA___ALGORITHM_COPY_H
#define __CUDA___ALGORITHM_COPY_H

#include <cuda/std/detail/__config>

#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

#if _CCCL_HAS_CTK() && !_CCCL_COMPILER(NVRTC)

#  include <cuda/__algorithm/common.h>
#  include <cuda/__stream/launch_transform.h>
#  include <cuda/__stream/stream_ref.h>
#  include <cuda/__type_traits/is_trivially_copyable.h>
#  include <cuda/std/__concepts/concept_macros.h>
#  include <cuda/std/__exception/exception_macros.h>
#  include <cuda/std/__host_stdlib/stdexcept>
#  include <cuda/std/mdspan>
#  include <cuda/std/span>

#  include <cuda/std/__cccl/prologue.h>

_CCCL_BEGIN_NAMESPACE_CUDA

//! @brief Source access order for copy_bytes
enum class source_access_order
{
  //! @brief Access source in stream order
  stream = 0x1,
  //! @brief Access source during the copy call, source can be destroyed after the API returns
  during_api_call = 0x2,
  //! @brief Access source in any order, the order can change across CUDA releases
  any = 0x3,
};

#  if _CCCL_CTK_AT_LEAST(13, 0)
static_assert(::cuda::std::to_underlying(source_access_order::stream) == ::cudaMemcpySrcAccessOrderStream);
static_assert(::cuda::std::to_underlying(source_access_order::during_api_call)
              == ::cudaMemcpySrcAccessOrderDuringApiCall);
static_assert(::cuda::std::to_underlying(source_access_order::any) == ::cudaMemcpySrcAccessOrderAny);
#  endif // _CCCL_CTK_AT_LEAST(13, 0)

//! @brief Configuration for copy_bytes
struct copy_configuration
{
  //! @brief Source memory location hint for copy_bytes, used only for managed memory
  memory_location src_location_hint = {};
  //! @brief Destination memory location hint for copy_bytes, used only for managed memory
  memory_location dst_location_hint = {};
  //! @brief Source access order for copy_bytes
  source_access_order src_access_order = source_access_order::any;
};

namespace __detail
{
template <class _SrcTp, class _DstTp>
_CCCL_CONCEPT __can_span_copy_bytes = _CCCL_REQUIRES_EXPR((_SrcTp, _DstTp), )(
  requires((!::cuda::std::is_const_v<_DstTp>) ),
  requires(::cuda::is_trivially_copyable_v<_SrcTp>),
  requires(::cuda::is_trivially_copyable_v<_DstTp>));

template <typename _SrcTy, ::cuda::std::size_t _SrcSize, typename _DstTy, ::cuda::std::size_t _DstSize>
_CCCL_HOST_API void __copy_bytes_impl(
  stream_ref __stream,
  ::cuda::std::span<_SrcTy, _SrcSize> __src,
  ::cuda::std::span<_DstTy, _DstSize> __dst,
  [[maybe_unused]] copy_configuration __config)
{
  static_assert(__can_span_copy_bytes<_SrcTy, _DstTy>);

  // If neither are dynamic_extent then they are a number, and in that case we can check at compile-time
  if constexpr ((_SrcSize != ::cuda::std::dynamic_extent) && (_DstSize != ::cuda::std::dynamic_extent))
  {
    // Can't use size_bytes() here because while the sizes are statically determinable, __src
    // and __dest are not core constexpr expressions due to be function arguments.
    constexpr auto __src_size  = _SrcSize * sizeof(_SrcTy);
    constexpr auto __dest_size = _DstSize * sizeof(_DstTy);

    static_assert(__dest_size >= __src_size, "Copy destination is too small to fit the source data");
  }
  else
  {
    if (__src.size_bytes() > __dst.size_bytes())
    {
      _CCCL_THROW(::std::invalid_argument, "Copy destination is too small to fit the source data");
    }
  }

  if (__src.empty())
  {
    return;
  }

#  if _CCCL_CTK_AT_LEAST(13, 0)
  CUmemcpyAttributes __attributes = {};
  __attributes.srcAccessOrder     = static_cast<::CUmemcpySrcAccessOrder>(__config.src_access_order);
  __attributes.srcLocHint.id      = __config.src_location_hint.id;
  __attributes.srcLocHint.type    = static_cast<::CUmemLocationType>(__config.src_location_hint.type);
  __attributes.dstLocHint.id      = __config.dst_location_hint.id;
  __attributes.dstLocHint.type    = static_cast<::CUmemLocationType>(__config.dst_location_hint.type);

  ::cuda::__ensure_current_context guard(__stream);
  ::cuda::__driver::__memcpyAsyncWithAttributes(
    __dst.data(), __src.data(), __src.size_bytes(), __stream.get(), __attributes);
#  else
  ::cuda::__driver::__memcpyAsync(__dst.data(), __src.data(), __src.size_bytes(), __stream.get());
#  endif // _CCCL_CTK_BELOW(13, 0)
}

template <class _SrcTp, class _SrcExtents, class _SrcLayout, class _DstTp, class _DstExtents, class _DstLayout>
_CCCL_CONCEPT __can_mdspan_copy_bytes =
  _CCCL_REQUIRES_EXPR((_SrcTp, _SrcExtents, _SrcLayout, _DstTp, _DstExtents, _DstLayout), )(
    requires(__can_span_copy_bytes<_SrcTp, _DstTp>),
    requires(::cuda::std::is_constructible_v<_DstExtents, _SrcExtents>),
    requires(::cuda::std::is_same_v<_SrcLayout, _DstLayout>));

template <typename _SrcElem,
          typename _SrcExtents,
          typename _SrcLayout,
          typename _SrcAccessor,
          typename _DstElem,
          typename _DstExtents,
          typename _DstLayout,
          typename _DstAccessor>
_CCCL_HOST_API void __copy_bytes_impl(
  stream_ref __stream,
  ::cuda::std::mdspan<_SrcElem, _SrcExtents, _SrcLayout, _SrcAccessor> __src,
  ::cuda::std::mdspan<_DstElem, _DstExtents, _DstLayout, _DstAccessor> __dst,
  copy_configuration __config)
{
  static_assert(__can_mdspan_copy_bytes<_SrcElem, _SrcExtents, _SrcLayout, _DstElem, _DstExtents, _DstLayout>);

  // Check only destination, because the layout of destination is the same as source
  if (!__dst.is_exhaustive())
  {
    _CCCL_THROW(::std::invalid_argument, "copy_bytes supports only exhaustive mdspans");
  }

  if (__src.extents() != __dst.extents())
  {
    _CCCL_THROW(::std::invalid_argument, "Copy destination size differs from the source");
  }

  ::cuda::__detail::__copy_bytes_impl(
    __stream,
    ::cuda::std::span(__src.data_handle(), __src.mapping().required_span_size()),
    ::cuda::std::span(__dst.data_handle(), __dst.mapping().required_span_size()),
    __config);
}
} // namespace __detail

//! @brief Launches a bytewise memory copy from source to destination into the provided
//! stream.
//!
//! Both source and destination needs to be a `contiguous_range` and convert to
//! `cuda::std::span`. The element types of both the source and destination range is
//! required to be trivially copyable.
//!
//! This call might be synchronous if either source or destination is pagable host memory.
//! It will be synchronous if both destination and copy is located in host memory.
//!
//! @param __stream Stream that the copy should be inserted into
//! @param __src Source to copy from
//! @param __dst Destination to copy into
//! @param __config Configuration for the copy
_CCCL_TEMPLATE(typename _SrcTy, typename _DstTy)
_CCCL_REQUIRES( // NOLINT(modernize-type-traits)
  __spannable<transformed_device_argument_t<_SrcTy>> _CCCL_AND __spannable<transformed_device_argument_t<_DstTy>>)
_CCCL_HOST_API void copy_bytes(stream_ref __stream, _SrcTy&& __src, _DstTy&& __dst, copy_configuration __config = {})
{
  ::cuda::__detail::__copy_bytes_impl(
    __stream,
    ::cuda::std::span(launch_transform(__stream, ::cuda::std::forward<_SrcTy>(__src))),
    ::cuda::std::span(launch_transform(__stream, ::cuda::std::forward<_DstTy>(__dst))),
    __config);
}

//! @overload
//! @note This overload accepts mdspan-compatible types.
_CCCL_TEMPLATE(typename _SrcTy, typename _DstTy)
_CCCL_REQUIRES( // NOLINT(modernize-type-traits)
  __mdspannable<transformed_device_argument_t<_SrcTy>> _CCCL_AND __mdspannable<transformed_device_argument_t<_DstTy>>)
_CCCL_HOST_API void copy_bytes(stream_ref __stream, _SrcTy&& __src, _DstTy&& __dst, copy_configuration __config = {})
{
  ::cuda::__detail::__copy_bytes_impl(
    __stream,
    ::cuda::__as_mdspan(launch_transform(__stream, ::cuda::std::forward<_SrcTy>(__src))),
    ::cuda::__as_mdspan(launch_transform(__stream, ::cuda::std::forward<_DstTy>(__dst))),
    __config);
}

_CCCL_END_NAMESPACE_CUDA

#  include <cuda/std/__cccl/epilogue.h>

#endif // _CCCL_HAS_CTK() && !_CCCL_COMPILER(NVRTC)

#endif // __CUDA___ALGORITHM_COPY_H