// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include <cuda/__cccl_config>
#include <cuda/atomic>
#include <cuda/buffer>
#include <cuda/functional>
#include <cuda/iterator>
#include <cuda/std/algorithm>
#include <cuda/std/cstddef>
#include <cuda/std/cstdint>
#include <cuda/std/functional>
#include <cuda/std/limits>
#include <cuda/std/span>
#include <cuda/std/type_traits>
#include <cuda/stream>

#include <cuda/experimental/__cuco/fixed_capacity_set.cuh>
#include <cuda/experimental/__cuco/probing_scheme.cuh>

#include <testing.cuh>

#include "test_common.cuh"

using key_types     = c2h::type_list<cuda::std::uint8_t, cuda::std::uint16_t, cuda::std::int32_t, cuda::std::int64_t>;
using cg_sizes      = c2h::type_list<int_c<1>, int_c<2>>;
using bucket_sizes  = c2h::type_list<int_c<1>, int_c<2>>;
using probing_kinds = c2h::type_list<int_c<0>, int_c<1>>;

template <class Key>
struct duplicate_key
{
  [[nodiscard]] _CCCL_HOST_DEVICE_API Key operator()(int i) const noexcept
  {
    return static_cast<Key>(i % 17);
  }
};

struct is_odd
{
  [[nodiscard]] _CCCL_DEVICE_API bool operator()(int value) const noexcept
  {
    return value % 2 != 0;
  }
};

struct matches_even_membership
{
  const int* results;
  int num_present;

  [[nodiscard]] _CCCL_DEVICE_API bool operator()(int i) const noexcept
  {
    return results[i] == static_cast<int>(i < num_present && i % 2 == 0);
  }
};

// Counts the callbacks per stored key through a non-const call operator. With `stride == 2`, only
// even keys are counted, at index `key / 2`, so a callback given an odd query key is not counted.
template <class Key>
struct record_visit
{
  int* visits{};
  int stride = 1;
  int calls  = 0;

  _CCCL_DEVICE_API void operator()(Key key) noexcept
  {
    ++calls;
    const int value = static_cast<int>(key);
    if (value % stride == 0)
    {
      cuda::atomic_ref<int, cuda::thread_scope_device>{visits[value / stride]}.fetch_add(1, cuda::memory_order_relaxed);
    }
  }
};

template <class Key>
struct is_below
{
  int bound;

  [[nodiscard]] _CCCL_DEVICE_API bool operator()(Key key) const noexcept
  {
    return static_cast<int>(key) < bound;
  }
};

struct is_negative
{
  [[nodiscard]] _CCCL_DEVICE_API bool operator()(int value) const noexcept
  {
    return value < 0;
  }
};

C2H_TEST("fixed_capacity_set insert and contains", "[container]", key_types, cg_sizes, bucket_sizes, probing_kinds)
{
  using key_type                              = c2h::get<0, TestType>;
  [[maybe_unused]] constexpr int cg_size      = c2h::get<1, TestType>::value;
  [[maybe_unused]] constexpr int bucket_size  = c2h::get<2, TestType>::value;
  [[maybe_unused]] constexpr int probing_kind = c2h::get<3, TestType>::value;
  using probing_type =
    cuda::std::conditional_t<probing_kind == 0,
                             cudax::cuco::linear_probing<cg_size, cuda::hash<key_type>>,
                             cudax::cuco::double_hashing<cg_size, cuda::hash<key_type>>>;
  using set_type = cudax::cuco::fixed_capacity_set<
    key_type,
    cuda::std::dynamic_extent,
    cuda::thread_scope_device,
    cuda::std::equal_to<key_type>,
    probing_type,
    bucket_size>;

  // Both present and absent query keys fit the smallest key type, with room for sentinels.
  constexpr int num_keys    = 41;
  constexpr int num_queries = 2 * num_keys;
  test_context context;
  set_type set{
    context.stream, context.mr, cuda::std::size_t{num_keys} * 2, cudax::cuco::empty_key{static_cast<key_type>(-1)}};
  const auto first = cuda::counting_iterator<key_type>{0};
  // Offset from the keys so that `pred(query_stencil[i])` and `pred(first[i])` disagree.
  const auto query_stencil = cuda::counting_iterator<int>{1};
  auto found               = cuda::make_buffer<int>(context.stream, context.mr, num_queries, 7);

  REQUIRE(set.insert(context.stream, first, first) == 0);
  set.insert_async(context.stream, first, first);
  set.contains(context.stream, first, first, found.begin());
  set.contains_async(context.stream, first, first, found.begin());
  set.contains_if(context.stream, first, first, query_stencil, is_odd{}, found.begin());
  set.contains_if_async(context.stream, first, first, query_stencil, is_odd{}, found.begin());
  REQUIRE(context.all_equal(found.data(), num_queries, 7));

  set.contains(context.stream, first, first + num_queries, found.begin());
  REQUIRE(context.matches(found.data(), num_queries, 0));
  REQUIRE(set.insert(context.stream, first, first + num_keys / 2) == num_keys / 2);
  set.insert_async(context.stream, first + num_keys / 2, first + num_keys);
  set.contains_async(context.stream, first, first + num_queries, found.begin());
  REQUIRE(context.matches(found.data(), num_queries, num_keys));
  set.contains_if(context.stream, first, first + num_queries, query_stencil, is_odd{}, found.begin());
  REQUIRE(cuda::std::all_of(
    context.policy(),
    cuda::counting_iterator<int>{0},
    cuda::counting_iterator<int>{num_queries},
    matches_even_membership{found.data(), num_keys}));
  REQUIRE(set.insert(context.stream, first, first + num_keys) == 0);
  auto retrieved = cuda::make_buffer<key_type>(context.stream, context.mr, num_queries, cuda::no_init);
  // Stored keys are distinct, so `num_keys` keys below `num_keys` are exactly the inserted keys.
  const auto retrieved_end = set.retrieve_all(context.stream, retrieved.data());
  REQUIRE(retrieved_end - retrieved.data() == num_keys);
  REQUIRE(cuda::std::all_of(context.policy(), retrieved.data(), retrieved_end, is_below<key_type>{num_keys}));

  // Each present key is visited exactly once; absent keys and empty ranges invoke no callback.
  auto visits = cuda::make_buffer<int>(context.stream, context.mr, num_queries, 0);
  set.for_each(context.stream, first, first, record_visit<key_type>{visits.data()});
  set.for_each_async(context.stream, first, first, record_visit<key_type>{visits.data()});
  set.for_each_async(context.stream, first, first + num_queries, record_visit<key_type>{visits.data()});
  REQUIRE(context.matches(visits.data(), num_queries, num_keys));
  set.contains(context.stream, first, first + num_queries, cuda::discard_iterator{});

  set.clear_async(context.stream);
  set.contains(context.stream, first, first + num_queries, found.begin());
  REQUIRE(context.matches(found.data(), num_queries, 0));
  set.for_each(context.stream, first, first + num_queries, record_visit<key_type>{visits.data()});
  REQUIRE(context.matches(visits.data(), num_queries, num_keys));
  REQUIRE(set.retrieve_all(context.stream, retrieved.data()) == retrieved.data());

  // Duplicates within one launch must be counted once even under contention.
  const auto duplicates = cuda::transform_iterator{cuda::counting_iterator<int>{0}, duplicate_key<key_type>{}};
  REQUIRE(set.insert(context.stream, duplicates, duplicates + 400) == 17);
  set.contains(context.stream, first, first + num_queries, found.begin());
  REQUIRE(context.matches(found.data(), num_queries, 17));

  set.clear(context.stream);
  set.insert_async(context.stream, first, first + num_keys);
  set.contains(context.stream, first, first + num_queries, found.begin());
  REQUIRE(context.matches(found.data(), num_queries, num_keys));

  // The stencil is negative only for the first half while every key is nonnegative, so the stencil, not the key,
  // selects the first half.
  const auto stencil = cuda::counting_iterator<int>{-num_keys / 2};
  set.clear(context.stream);
  REQUIRE(set.insert_if(context.stream, first, first + num_keys, stencil, is_negative{}) == num_keys / 2);
  set.contains(context.stream, first, first + num_queries, found.begin());
  REQUIRE(context.matches(found.data(), num_queries, num_keys / 2));

  set.clear(context.stream);
  set.insert_if_async(context.stream, first, first + num_keys, stencil, is_negative{});
  set.contains(context.stream, first, first + num_queries, found.begin());
  REQUIRE(context.matches(found.data(), num_queries, num_keys / 2));

  set_type full{context.stream, context.mr, cuda::std::size_t{13}, cudax::cuco::empty_key{static_cast<key_type>(-1)}};
  const int capacity = static_cast<int>(full.capacity());
  REQUIRE(full.insert(context.stream, first, first + capacity) == full.capacity());
  REQUIRE(full.insert(context.stream, first, first + capacity) == 0);
  REQUIRE(full.insert(context.stream, first + capacity, first + capacity + 1) == 0);
  const auto full_end = full.retrieve_all(context.stream, retrieved.data());
  REQUIRE(full_end - retrieved.data() == capacity);
  REQUIRE(cuda::std::all_of(context.policy(), retrieved.data(), full_end, is_below<key_type>{capacity}));
  full.insert_async(context.stream, first + capacity, first + capacity + 1);
  full.contains(context.stream, first, first + capacity + 1, found.begin());
  REQUIRE(context.matches(found.data(), capacity + 1, capacity));
}

struct insert_key
{
  int value;

  [[nodiscard]] _CCCL_HOST_DEVICE_API explicit operator int() const noexcept
  {
    return value;
  }
};

// No conversion to the stored key type: lookup must use the supplied hash and equality.
struct query_key
{
  int value;
};

struct heterogeneous_hash
{
  cuda::std::uint32_t seed = 0;

  [[nodiscard]] _CCCL_HOST_DEVICE_API cuda::std::uint32_t operator()(int value) const noexcept
  {
    return static_cast<cuda::std::uint32_t>(value / 2) ^ seed;
  }

  template <class Key>
  [[nodiscard]] _CCCL_HOST_DEVICE_API cuda::std::uint32_t operator()(Key key) const noexcept
  {
    return operator()(key.value);
  }
};

struct heterogeneous_equal
{
  [[nodiscard]] _CCCL_HOST_DEVICE_API bool operator()(int lhs, int rhs) const noexcept
  {
    return lhs / 2 == rhs / 2;
  }

  template <class Key>
  [[nodiscard]] _CCCL_HOST_DEVICE_API bool operator()(Key lhs, int rhs) const noexcept
  {
    return operator()(lhs.value, rhs);
  }
};

template <class Key>
struct make_heterogeneous_key
{
  int offset;

  [[nodiscard]] _CCCL_HOST_DEVICE_API Key operator()(int i) const noexcept
  {
    return Key{(2 * i) + offset};
  }
};

C2H_TEST(
  "fixed_capacity_set heterogeneous keys and custom equivalence", "[container][heterogeneous]", cg_sizes, probing_kinds)
{
  [[maybe_unused]] constexpr int cg_size      = c2h::get<0, TestType>::value;
  [[maybe_unused]] constexpr int probing_kind = c2h::get<1, TestType>::value;
  using probing_type =
    cuda::std::conditional_t<probing_kind == 0,
                             cudax::cuco::linear_probing<cg_size, heterogeneous_hash>,
                             cudax::cuco::double_hashing<cg_size, heterogeneous_hash>>;
  using set_type = cudax::cuco::
    fixed_capacity_set<int, cuda::std::dynamic_extent, cuda::thread_scope_device, heterogeneous_equal, probing_type>;
  constexpr int num_keys = 37;
  test_context context;
  set_type set{context.stream, context.mr, cuda::std::size_t{100}, cudax::cuco::empty_key{-100}, {}, probing_type{}};
  const auto first      = cuda::counting_iterator<int>{0};
  const auto input      = cuda::transform_iterator{first, make_heterogeneous_key<insert_key>{0}};
  const auto equivalent = cuda::transform_iterator{first, make_heterogeneous_key<insert_key>{1}};
  const auto queries    = cuda::transform_iterator{first, make_heterogeneous_key<query_key>{1}};
  auto found            = cuda::make_buffer<int>(context.stream, context.mr, 2 * num_keys, -1);

  set.contains(context.stream, queries, queries + 2 * num_keys, found.begin());
  REQUIRE(context.matches(found.data(), 2 * num_keys, 0));
  REQUIRE(set.insert(context.stream, input, input + num_keys) == num_keys);
  REQUIRE(set.insert(context.stream, equivalent, equivalent + num_keys) == 0);
  set.contains_async(context.stream, queries, queries + 2 * num_keys, found.begin());
  REQUIRE(context.matches(found.data(), 2 * num_keys, num_keys));

  // The callback receives the stored even key, not the odd query.
  auto visits = cuda::make_buffer<int>(context.stream, context.mr, 2 * num_keys, 0);
  set.for_each(context.stream, queries, queries + 2 * num_keys, record_visit<int>{visits.data(), 2});
  REQUIRE(context.matches(visits.data(), 2 * num_keys, num_keys));
}

C2H_TEST("fixed_capacity_set operations on another stream", "[container][stream]", key_types)
{
  using key_type         = c2h::get<0, TestType>;
  constexpr int num_keys = 100;
  const auto first       = cuda::counting_iterator<key_type>{0};
  test_context context;
  const cuda::stream other{context.stream.device()};
  cudax::cuco::fixed_capacity_set<key_type> set{
    context.stream, context.mr, num_keys * 2, cudax::cuco::empty_key{static_cast<key_type>(-1)}};
  auto found = cuda::make_buffer<int>(other, context.mr, num_keys, 0);
  // CCCL construction is asynchronous: establish the dependency before switching streams.
  context.stream.sync();
  REQUIRE(set.insert(other, first, first + num_keys) == num_keys);
  set.contains(other, first, first + num_keys, found.begin());
  REQUIRE(context.all_equal(found.data(), num_keys, 1));

  set.clear(context.stream);
  REQUIRE(set.insert(context.stream, first, first + num_keys) == num_keys);
  set.contains(other, first, first + num_keys, found.begin());
  REQUIRE(context.all_equal(found.data(), num_keys, 1));
}

struct hashed_key
{
  cuda::std::uint32_t hash;
  cuda::std::int32_t value;
};

struct stored_hash
{
  [[nodiscard]] _CCCL_HOST_DEVICE_API cuda::std::uint32_t operator()(hashed_key key) const noexcept
  {
    return key.hash;
  }
};

struct distinct_keys
{
  [[nodiscard]] _CCCL_HOST_DEVICE_API bool operator()(hashed_key, hashed_key) const noexcept
  {
    return false;
  }
};

struct make_hashed_key
{
  [[nodiscard]] _CCCL_HOST_DEVICE_API hashed_key operator()(cuda::std::int32_t value) const noexcept
  {
    return {static_cast<cuda::std::uint32_t>(cuda::hash<cuda::std::int32_t>{}(value)), value};
  }
};

C2H_TEST("fixed_capacity_set custom key atomic storage regression", "[container][atomic]")
{
  // cuCollections regression for spark-rapids#12586 / cudf#18587: a key storing its
  // hash must remain valid when copied through the packed atomic representation.
  using set_type =
    cudax::cuco::fixed_capacity_set<hashed_key,
                                    cuda::std::dynamic_extent,
                                    cuda::thread_scope_device,
                                    distinct_keys,
                                    cudax::cuco::linear_probing<1, stored_hash>>;
  constexpr int num_keys = 100'000;
  const test_context context{};
  set_type set{context.stream,
               context.mr,
               num_keys,
               0.5,
               cudax::cuco::empty_key{hashed_key{cuda::std::numeric_limits<cuda::std::uint32_t>::max(), -1}}};
  const auto first = cuda::transform_iterator{cuda::counting_iterator<cuda::std::int32_t>{0}, make_hashed_key{}};
  REQUIRE(set.insert(context.stream, first, first + num_keys) == num_keys);
}

C2H_TEST("fixed_capacity_set large input regression", "[container][.large]", cg_sizes)
{
  // The upstream 1.2B-key regression needs about 30 GiB. Opt in with '[.large]'.
  using key_type        = cuda::std::int64_t;
  constexpr int cg_size = c2h::get<0, TestType>::value;
  using set_type =
    cudax::cuco::fixed_capacity_set<key_type,
                                    cuda::std::dynamic_extent,
                                    cuda::thread_scope_device,
                                    cuda::std::equal_to<key_type>,
                                    cudax::cuco::double_hashing<cg_size, cuda::hash<key_type>>>;
  constexpr cuda::std::size_t num_keys = 1'200'000'000;
  test_context context;
  set_type set{context.stream, context.mr, num_keys * 2, cudax::cuco::empty_key<key_type>{-1}};
  const auto first = cuda::counting_iterator<key_type>{0};
  auto found       = cuda::make_buffer<bool>(context.stream, context.mr, num_keys, cuda::no_init);
  set.contains(context.stream, first, first + num_keys, found.begin());
  REQUIRE(cuda::std::none_of(context.policy(), found.begin(), found.end(), cuda::std::identity{}));
  REQUIRE(set.insert(context.stream, first, first + num_keys) == num_keys);
  set.contains(context.stream, first, first + num_keys, found.begin());
  REQUIRE(cuda::std::all_of(context.policy(), found.begin(), found.end(), cuda::std::identity{}));
  auto retrieved           = cuda::make_buffer<key_type>(context.stream, context.mr, num_keys, cuda::no_init);
  const auto retrieved_end = set.retrieve_all(context.stream, retrieved.data());
  REQUIRE(static_cast<cuda::std::size_t>(retrieved_end - retrieved.data()) == num_keys);
  REQUIRE(cuda::std::all_of(
    context.policy(), retrieved.data(), retrieved_end, is_below<key_type>{static_cast<int>(num_keys)}));
}