// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright Contributors to the Kokkos project

#include <gtest/gtest.h>

#include <Kokkos_Macros.hpp>
#ifdef KOKKOS_ENABLE_EXPERIMENTAL_CXX20_MODULES
import kokkos.core;
import kokkos.core_impl;
#else
#include <Kokkos_Core.hpp>
#endif

#include <regex>
#include <limits>
#include <type_traits>

namespace {

TEST(TEST_CATEGORY, range_policy_runtime_parameters) {
  using Policy = Kokkos::RangePolicy<>;
  using Index  = Policy::index_type;

  Index work_begin = -5;
  Index work_end   = 15;
  Index chunk_size = 10;
  {
    Policy p(work_begin, work_end);
    ASSERT_EQ(p.begin(), work_begin);
    ASSERT_EQ(p.end(), work_end);
  }
  {
    Policy p(Kokkos::DefaultExecutionSpace(), work_begin, work_end);
    ASSERT_EQ(p.begin(), work_begin);
    ASSERT_EQ(p.end(), work_end);
  }
  {
    Policy p(work_begin, work_end, Kokkos::ChunkSize(chunk_size));
    ASSERT_EQ(p.begin(), work_begin);
    ASSERT_EQ(p.end(), work_end);
    ASSERT_EQ(p.chunk_size(), chunk_size);
  }
  {
    Policy p(Kokkos::DefaultExecutionSpace(), work_begin, work_end,
             Kokkos::ChunkSize(chunk_size));
    ASSERT_EQ(p.begin(), work_begin);
    ASSERT_EQ(p.end(), work_end);
    ASSERT_EQ(p.chunk_size(), chunk_size);
  }
  {
    Policy p;  // default-constructed
    ASSERT_EQ(p.begin(), Index(0));
    ASSERT_EQ(p.end(), Index(0));
    ASSERT_EQ(p.chunk_size(), Index(0));

    // copy-assigned
    p = Policy(work_begin, work_end, Kokkos::ChunkSize(chunk_size));
    ASSERT_EQ(p.begin(), work_begin);
    ASSERT_EQ(p.end(), work_end);
    ASSERT_EQ(p.chunk_size(), chunk_size);
  }
  {
    Policy p1(work_begin, work_end, Kokkos::ChunkSize(chunk_size));
    Policy p2(p1);  // copy-constructed
    ASSERT_EQ(p1.begin(), p2.begin());
    ASSERT_EQ(p1.end(), p2.end());
    ASSERT_EQ(p1.chunk_size(), p2.chunk_size());
  }
}

TEST(TEST_CATEGORY_DEATH, range_policy_invalid_bounds) {
  ::testing::FLAGS_gtest_death_test_style = "threadsafe";

  using Policy    = Kokkos::RangePolicy<TEST_EXECSPACE>;
  using ChunkSize = Kokkos::ChunkSize;

  std::string msg =
      "Kokkos::RangePolicy bounds error: The lower bound (100) is greater than "
      "the upper bound (90).\n";
  // escape the parentheses in the regex to match the error message
  msg = std::regex_replace(msg, std::regex("\\(|\\)"), "\\$&");
  ASSERT_DEATH({ (void)Policy(100, 90); }, msg);

  ASSERT_DEATH({ (void)Policy(TEST_EXECSPACE(), 100, 90, ChunkSize(10)); },
               msg);
}
TEST(TEST_CATEGORY_DEATH, range_policy_check_exceeding_max) {
  ::testing::FLAGS_gtest_death_test_style = "threadsafe";
  // Trigger due to exceeding a policy's range maximum
  using IntPolicy = Kokkos::RangePolicy<int>;

  long long const n_large = 9223372036854774771;
  [[maybe_unused]] std::string msg =
      "Kokkos::RangePolicy bound type error: an unsafe implicit conversion is "
      "performed";
  ASSERT_DEATH((void)IntPolicy(-4, n_large), msg);
}

TEST(TEST_CATEGORY_DEATH, range_policy_check_exceeding_min) {
  ::testing::FLAGS_gtest_death_test_style = "threadsafe";
  // Trigger due to exceeding a policy's range minimum
  using IntPolicy = Kokkos::RangePolicy<int>;

  long long const n_small = -9223372036854774771;
  [[maybe_unused]] std::string msg =
      "Kokkos::RangePolicy bound type error: an unsafe implicit conversion is "
      "performed";
  ASSERT_DEATH((void)IntPolicy(n_small, 4), msg);
}

struct W {  // round-trip conversion check for narrowing should "fire"
  W() : val_(1) {}
  W(int const* ptr) : val_(*ptr) {}
  // Deliberately discards the int value so W -> member_type -> W loses
  // information for val_ != 1 (exercises the round-trip check).
  W(int const) : val_(1) {}
  operator int() const { return val_; }
  int val_;
};

TEST(TEST_CATEGORY_DEATH, range_policy_round_trip_conversion_fires) {
  ::testing::FLAGS_gtest_death_test_style = "threadsafe";

  using Policy = Kokkos::RangePolicy<>;

  static_assert(std::is_convertible_v<W, Policy::index_type>);
  static_assert(std::is_convertible_v<Policy::index_type, W>);

  int const n = 5;
  [[maybe_unused]] std::string msg =
      "Kokkos::RangePolicy bound type error: an unsafe implicit conversion is "
      "performed";
  ASSERT_DEATH((void)Policy(0, W(&n)), msg);
}

struct B {  // round-trip conversion would not compile
  B(int const* ptr) : val_(*ptr) {}
  operator int() const { return val_; }

  int val_;
};

TEST(TEST_CATEGORY, range_policy_one_way_convertible_bounds) {
  using Policy    = Kokkos::RangePolicy<>;
  using IndexType = Policy::index_type;

  static_assert(std::is_convertible_v<B, IndexType>);
  static_assert(!std::is_convertible_v<IndexType, B>);

  int const n = 1;
  Policy policy(0, B(&n));
  EXPECT_EQ(policy.begin(), static_cast<IndexType>(0));
  EXPECT_EQ(policy.end(), static_cast<IndexType>(1));
}

TEST(TEST_CATEGORY_DEATH, range_policy_check_sign_changes) {
  ::testing::FLAGS_gtest_death_test_style = "threadsafe";

  using UInt32Policy =
      Kokkos::RangePolicy<TEST_EXECSPACE, Kokkos::IndexType<std::uint32_t>>;

  [[maybe_unused]] std::string msg =
      "Kokkos::RangePolicy bound type error: an unsafe implicit conversion is "
      "performed";
  {
    std::int64_t n = std::numeric_limits<std::int64_t>::max();
    ASSERT_DEATH((void)UInt32Policy(0, n), msg);
  }
  {
    std::int64_t n = std::numeric_limits<std::int64_t>::min();
    ASSERT_DEATH((void)UInt32Policy(n, 0), msg);
  }
}

TEST(TEST_CATEGORY_DEATH, range_policy_implicitly_converted_bounds) {
  using UIntIndexType     = Kokkos::IndexType<unsigned>;
  using IntIndexType      = Kokkos::IndexType<int>;
  using UIntPolicy        = Kokkos::RangePolicy<TEST_EXECSPACE, UIntIndexType>;
  using IntPolicy         = Kokkos::RangePolicy<TEST_EXECSPACE, IntIndexType>;
  using DefaultExecPolicy = Kokkos::RangePolicy<Kokkos::DefaultExecutionSpace>;

  std::string msg =
      "Kokkos::RangePolicy bound type error: an unsafe implicit conversion is "
      "performed on a bound (), which may not preserve its original value.\n";

  [[maybe_unused]] auto get_error_msg = [](auto str, auto val) {
    return str.insert(str.find("(") + 1, std::to_string(val).c_str());
  };
  ::testing::FLAGS_gtest_death_test_style = "threadsafe";

  std::string expected = std::regex_replace(msg, std::regex("\\(|\\)"), "\\$&");
  {
    int test_val = -1;
    ASSERT_DEATH({ (void)UIntPolicy(test_val, 10); },
                 get_error_msg(expected, test_val));
  }
  {
    unsigned test_val = std::numeric_limits<unsigned>::max();
    ASSERT_DEATH({ (void)IntPolicy(0, test_val); },
                 get_error_msg(expected, test_val));
  }
  {
    unsigned long long test_val =
        std::numeric_limits<unsigned long long>::max();
    ASSERT_DEATH({ (void)IntPolicy(0, test_val); },
                 get_error_msg(expected, test_val));

    ASSERT_DEATH({ (void)DefaultExecPolicy(0, test_val); },
                 get_error_msg(expected, test_val));
  }
  {
    int test_val = -1;
    ASSERT_DEATH({ (void)UIntPolicy(test_val, 10, Kokkos::ChunkSize(2)); },
                 get_error_msg(expected, test_val));
  }
}

constexpr bool test_chunk_size_explicit() {
  using ExecutionSpace = TEST_EXECSPACE;
  using Kokkos::ChunkSize;

  static_assert(!std::is_convertible_v<int, ChunkSize>);
  static_assert(std::is_constructible_v<ChunkSize, int>);
  static_assert(
      !std::is_constructible_v<
          Kokkos::RangePolicy<Kokkos::DefaultExecutionSpace>, int, int, int>);
  static_assert(std::is_constructible_v<
                Kokkos::RangePolicy<Kokkos::DefaultExecutionSpace>, int, int,
                ChunkSize>);
  static_assert(!std::is_constructible_v<Kokkos::RangePolicy<ExecutionSpace>,
                                         ExecutionSpace, int, int, int>);
  static_assert(std::is_constructible_v<Kokkos::RangePolicy<ExecutionSpace>,
                                        ExecutionSpace, int, int, ChunkSize>);
  return true;
}

static_assert(test_chunk_size_explicit());

// The execution space is defaulted if not given to the constructor.
TEST(TEST_CATEGORY, range_policy_default_space) {
  using policy_t = Kokkos::RangePolicy<TEST_EXECSPACE>;

  policy_t defaulted(42, 666);

  ASSERT_EQ(defaulted.space(), TEST_EXECSPACE{});
}

// The execution space instance can be updated.
TEST(TEST_CATEGORY, range_policy_impl_set_space) {
  using policy_t = Kokkos::RangePolicy<TEST_EXECSPACE>;

  const auto [exec_old, exec_new] =
      Kokkos::Experimental::partition_space(TEST_EXECSPACE{}, 1, 1);

  const policy_t policy_old(exec_old, 42, 666);
  ASSERT_EQ(policy_old.space(), exec_old);

  const policy_t policy_new(Kokkos::Impl::PolicyUpdate{}, policy_old, exec_new);
  ASSERT_EQ(policy_new.space(), exec_new);
  ASSERT_EQ(policy_new.begin(), 42u);
  ASSERT_EQ(policy_new.end(), 666u);
}

}  // namespace
