// SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: BSD-3
#pragma once

#include <thrust/detail/config.h>

#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_CUDA_COMPILATION()

#  include <cub/device/device_merge.cuh>

#  include <thrust/detail/temporary_array.h>
#  include <thrust/iterator/iterator_traits.h>
#  include <thrust/system/cuda/detail/dispatch.h>
#  include <thrust/system/cuda/detail/util.h>

#  include <cuda/std/__functional/operations.h>
#  include <cuda/std/__iterator/distance.h>
#  include <cuda/std/__utility/pair.h>
#  include <cuda/std/cstdint>

THRUST_NAMESPACE_BEGIN
namespace cuda_cub
{
_CCCL_EXEC_CHECK_DISABLE
template <class Derived, class KeysIt1, class KeysIt2, class ResultIt, class CompareOp = ::cuda::std::less<>>
ResultIt _CCCL_HOST_DEVICE
merge(execution_policy<Derived>& policy,
      KeysIt1 keys1_begin,
      KeysIt1 keys1_end,
      KeysIt2 keys2_begin,
      KeysIt2 keys2_end,
      ResultIt result_begin,
      CompareOp compare_op = {})
{
  THRUST_CDP_DISPATCH(
    ({
      using size_type         = thrust::detail::it_difference_t<KeysIt1>;
      const auto num_keys1    = static_cast<size_type>(::cuda::std::distance(keys1_begin, keys1_end));
      const auto num_keys2    = static_cast<size_type>(::cuda::std::distance(keys2_begin, keys2_end));
      const auto num_keys_out = num_keys1 + num_keys2;
      if (num_keys_out == 0)
      {
        return result_begin;
      }

      const auto stream   = cuda_cub::stream(policy);
      size_t storage_size = 0;
      cudaError_t status  = cub::DeviceMerge::MergeKeys(
        nullptr, storage_size, keys1_begin, num_keys1, keys2_begin, num_keys2, result_begin, compare_op, stream);
      throw_on_error(status, "merge: failed on 1st step");

      thrust::detail::temporary_array<char, Derived> temp_storage(policy, storage_size);
      status = cub::DeviceMerge::MergeKeys(
        temp_storage.data().get(),
        storage_size,
        keys1_begin,
        num_keys1,
        keys2_begin,
        num_keys2,
        result_begin,
        compare_op,
        stream);
      throw_on_error(status, "merge: failed on 2nd step");

      status = cuda_cub::synchronize_optional(policy);
      throw_on_error(status, "merge: failed to synchronize");

      return result_begin + num_keys_out;
    }),
    ({
      return thrust::merge(
        cvt_to_seq(derived_cast(policy)), keys1_begin, keys1_end, keys2_begin, keys2_end, result_begin, compare_op);
    }));
}

_CCCL_EXEC_CHECK_DISABLE
template <class Derived,
          class KeysIt1,
          class KeysIt2,
          class ItemsIt1,
          class ItemsIt2,
          class KeysOutputIt,
          class ItemsOutputIt,
          class CompareOp = ::cuda::std::less<>>
::cuda::std::pair<KeysOutputIt, ItemsOutputIt> _CCCL_HOST_DEVICE merge_by_key(
  execution_policy<Derived>& policy,
  KeysIt1 keys1_begin,
  KeysIt1 keys1_end,
  KeysIt2 keys2_begin,
  KeysIt2 keys2_end,
  ItemsIt1 items1_begin,
  ItemsIt2 items2_begin,
  KeysOutputIt keys_out_begin,
  ItemsOutputIt items_out_begin,
  CompareOp compare_op = {})
{
  THRUST_CDP_DISPATCH(({
                        using size_type = thrust::detail::it_difference_t<KeysIt1>;

                        const auto num_keys1    = static_cast<size_type>(::cuda::std::distance(keys1_begin, keys1_end));
                        const auto num_keys2    = static_cast<size_type>(::cuda::std::distance(keys2_begin, keys2_end));
                        const auto num_keys_out = num_keys1 + num_keys2;
                        if (num_keys_out == 0)
                        {
                          return {keys_out_begin, items_out_begin};
                        }

                        const auto stream   = cuda_cub::stream(policy);
                        size_t storage_size = 0;
                        cudaError_t status  = cub::DeviceMerge::MergePairs(
                          nullptr,
                          storage_size,
                          keys1_begin,
                          items1_begin,
                          num_keys1,
                          keys2_begin,
                          items2_begin,
                          num_keys2,
                          keys_out_begin,
                          items_out_begin,
                          compare_op,
                          stream);
                        throw_on_error(status, "merge: failed on 1st step");

                        thrust::detail::temporary_array<char, Derived> temp_storage(policy, storage_size);
                        status = cub::DeviceMerge::MergePairs(
                          temp_storage.data().get(),
                          storage_size,
                          keys1_begin,
                          items1_begin,
                          num_keys1,
                          keys2_begin,
                          items2_begin,
                          num_keys2,
                          keys_out_begin,
                          items_out_begin,
                          compare_op,
                          stream);
                        throw_on_error(status, "merge: failed on 2nd step");

                        status = cuda_cub::synchronize_optional(policy);
                        throw_on_error(status, "merge: failed to synchronize");

                        return {keys_out_begin + num_keys_out, items_out_begin + num_keys_out};
                      }),
                      ({
                        return thrust::merge_by_key(
                          cvt_to_seq(derived_cast(policy)),
                          keys1_begin,
                          keys1_end,
                          keys2_begin,
                          keys2_end,
                          items1_begin,
                          items2_begin,
                          keys_out_begin,
                          items_out_begin,
                          compare_op);
                      }));
}
} // namespace cuda_cub
THRUST_NAMESPACE_END
#endif // _CCCL_CUDA_COMPILATION()
