#include <thrust/fill.h>
#include <thrust/logical.h>
#include <thrust/memory.h>
#include <thrust/reverse.h>
#include <thrust/sequence.h>
#include <thrust/sort.h>

#include <cuda/std/utility>

#include <iostream>

#include <unittest/unittest.h>

// WAR NVIDIA/cccl#1731
// Some tests miscompile for non-CUDA backends on MSVC 2017 and 2019 (though 2022 is fine).
// This is due to a bug in the compiler that breaks __THRUST_DEFINE_HAS_MEMBER_FUNCTION.
#if defined(_MSC_VER) && _MSC_VER <= 1929 && THRUST_DEVICE_SYSTEM != THRUST_DEVICE_SYSTEM_CUDA
#  define WAR_BUG_1731
#endif

// Define a new system class, as the my_system one is already used with a thrust::sort template definition
// that calls back into sort.cu
class my_memory_system : public thrust::device_execution_policy<my_memory_system>
{
public:
  // disallow default construction
  my_memory_system() = delete;

  my_memory_system(int)
      : correctly_dispatched(false)
      , num_copies(0)
  {}

  my_memory_system(const my_memory_system& other)
      : correctly_dispatched(false)
      , num_copies(other.num_copies + 1)
  {}

  void validate_dispatch()
  {
    correctly_dispatched = (num_copies == 0);
  }

  bool is_valid()
  {
    return correctly_dispatched;
  }

private:
  bool correctly_dispatched;

  // count the number of copies so that we can validate
  // that dispatch does not introduce any
  unsigned int num_copies;
};

namespace my_old_namespace
{
struct my_old_temporary_allocation_system : public thrust::device_execution_policy<my_old_temporary_allocation_system>
{};

template <typename T>
cuda::std::pair<thrust::pointer<T, my_old_temporary_allocation_system>, std::ptrdiff_t>
get_temporary_buffer(my_old_temporary_allocation_system, std::ptrdiff_t)
{
  thrust::pointer<T, my_old_temporary_allocation_system> const result(reinterpret_cast<T*>(4217));

  return cuda::std::make_pair(result, 314);
}

template <typename Pointer>
void return_temporary_buffer(my_old_temporary_allocation_system, Pointer p)
{
  using RP = typename cuda::std::pointer_traits<Pointer>::raw_pointer;
  ASSERT_EQUAL(p.get(), reinterpret_cast<RP>(4217));
}
} // namespace my_old_namespace

namespace my_new_namespace
{
struct my_new_temporary_allocation_system : public thrust::device_execution_policy<my_new_temporary_allocation_system>
{};

template <typename T>
cuda::std::pair<thrust::pointer<T, my_new_temporary_allocation_system>, std::ptrdiff_t>
get_temporary_buffer(my_new_temporary_allocation_system, std::ptrdiff_t)
{
  thrust::pointer<T, my_new_temporary_allocation_system> const result(reinterpret_cast<T*>(1742));

  return cuda::std::make_pair(result, 413);
}

template <typename Pointer>
void return_temporary_buffer(my_new_temporary_allocation_system, Pointer)
{
  // This should never be called (the three-argument with size overload below
  // should be preferred) and shouldn't be ambiguous.
  ASSERT_EQUAL(true, false);
}

template <typename Pointer>
void return_temporary_buffer(my_new_temporary_allocation_system, Pointer p, std::ptrdiff_t n)
{
  using RP = typename cuda::std::pointer_traits<Pointer>::raw_pointer;
  ASSERT_EQUAL(p.get(), reinterpret_cast<RP>(1742));
  ASSERT_EQUAL(n, 413);
}
} // namespace my_new_namespace

template <typename T1, typename T2>
bool are_same(const T1&, const T2&)
{
  return false;
}

template <typename T>
bool are_same(const T&, const T&)
{
  return true;
}

void TestSelectSystemDifferentTypes()
{
  using thrust::system::detail::generic::select_system;

  my_memory_system my_sys(0);
  thrust::device_system_tag device_sys;

  // select_system(my_system, device_system_tag) should return device_system_tag (the minimum tag)
  bool is_device_system_tag = are_same(device_sys, select_system(my_sys, device_sys));
  ASSERT_EQUAL(true, is_device_system_tag);

  // select_system(device_system_tag, my_tag) should return device_system_tag (the minimum tag)
  is_device_system_tag = are_same(device_sys, select_system(device_sys, my_sys));
  ASSERT_EQUAL(true, is_device_system_tag);
}
DECLARE_UNITTEST(TestSelectSystemDifferentTypes);

void TestSelectSystemSameTypes()
{
  using thrust::system::detail::generic::select_system;

  my_memory_system my_sys(0);
  thrust::device_system_tag device_sys;
  thrust::host_system_tag host_sys;

  // select_system(host_system_tag, host_system_tag) should return host_system_tag
  bool is_host_system_tag = are_same(host_sys, select_system(host_sys, host_sys));
  ASSERT_EQUAL(true, is_host_system_tag);

  // select_system(device_system_tag, device_system_tag) should return device_system_tag
  bool is_device_system_tag = are_same(device_sys, select_system(device_sys, device_sys));
  ASSERT_EQUAL(true, is_device_system_tag);

  // select_system(my_system, my_system) should return my_system
  bool is_my_system = are_same(my_sys, select_system(my_sys, my_sys));
  ASSERT_EQUAL(true, is_my_system);
}
DECLARE_UNITTEST(TestSelectSystemSameTypes);

void TestGetTemporaryBuffer()
{
  const std::ptrdiff_t n = 9001;

  thrust::device_system_tag dev_tag;
  using pointer                                       = thrust::pointer<int, thrust::device_system_tag>;
  cuda::std::pair<pointer, std::ptrdiff_t> ptr_and_sz = thrust::get_temporary_buffer<int>(dev_tag, n);

  ASSERT_EQUAL(ptr_and_sz.second, n);

  const int ref_val = 13;
  thrust::device_vector<int> ref(n, ref_val);

  thrust::fill_n(ptr_and_sz.first, n, ref_val);

  ASSERT_EQUAL(true, thrust::all_of(ptr_and_sz.first, ptr_and_sz.first + n, thrust::placeholders::_1 == ref_val));

  thrust::return_temporary_buffer(dev_tag, ptr_and_sz.first, ptr_and_sz.second);
}
DECLARE_UNITTEST(TestGetTemporaryBuffer);

void TestMalloc()
{
  const std::ptrdiff_t n = 9001;

  thrust::device_system_tag dev_tag;
  using pointer = thrust::pointer<int, thrust::device_system_tag>;
  pointer ptr   = pointer(static_cast<int*>(thrust::malloc(dev_tag, sizeof(int) * n).get()));

  const int ref_val = 13;
  thrust::device_vector<int> ref(n, ref_val);

  thrust::fill_n(ptr, n, ref_val);

  ASSERT_EQUAL(true, thrust::all_of(ptr, ptr + n, thrust::placeholders::_1 == ref_val));

  thrust::free(dev_tag, ptr);
}
DECLARE_UNITTEST(TestMalloc);

thrust::pointer<void, my_memory_system> malloc(my_memory_system& system, std::size_t)
{
  system.validate_dispatch();

  return thrust::pointer<void, my_memory_system>();
}

void TestMallocDispatchExplicit()
{
  const size_t n = 0;

  my_memory_system sys(0);
  thrust::malloc(sys, n);

  ASSERT_EQUAL(true, sys.is_valid());
}
DECLARE_UNITTEST(TestMallocDispatchExplicit);

template <typename Pointer>
void free(my_memory_system& system, Pointer)
{
  system.validate_dispatch();
}

void TestFreeDispatchExplicit()
{
  thrust::pointer<my_memory_system, void> ptr;

  my_memory_system sys(0);
  thrust::free(sys, ptr);

  ASSERT_EQUAL(true, sys.is_valid());
}
DECLARE_UNITTEST(TestFreeDispatchExplicit);

template <typename T>
cuda::std::pair<thrust::pointer<T, my_memory_system>, std::ptrdiff_t>
get_temporary_buffer(my_memory_system& system, std::ptrdiff_t n)
{
  system.validate_dispatch();

  thrust::device_system_tag device_sys;
  cuda::std::pair<thrust::pointer<T, thrust::device_system_tag>, std::ptrdiff_t> result =
    thrust::get_temporary_buffer<T>(device_sys, n);
  return cuda::std::make_pair(thrust::pointer<T, my_memory_system>(result.first.get()), result.second);
}

void TestGetTemporaryBufferDispatchExplicit()
{
  const std::ptrdiff_t n = 9001;

  my_memory_system sys(0);
  using pointer                                       = thrust::pointer<int, thrust::device_system_tag>;
  cuda::std::pair<pointer, std::ptrdiff_t> ptr_and_sz = thrust::get_temporary_buffer<int>(sys, n);

  ASSERT_EQUAL(ptr_and_sz.second, n);
  ASSERT_EQUAL(true, sys.is_valid());

  const int ref_val = 13;
  thrust::device_vector<int> ref(n, ref_val);

  thrust::fill_n(ptr_and_sz.first, n, ref_val);

  ASSERT_EQUAL(true, thrust::all_of(ptr_and_sz.first, ptr_and_sz.first + n, thrust::placeholders::_1 == ref_val));

  thrust::return_temporary_buffer(sys, ptr_and_sz.first, ptr_and_sz.second);
}
DECLARE_UNITTEST(TestGetTemporaryBufferDispatchExplicit);

#ifndef WAR_BUG_1731

void TestGetTemporaryBufferDispatchImplicit()
{
  if (are_same(thrust::device_system_tag(), thrust::system::cpp::tag()))
  {
    // XXX cpp uses the internal scalar backend, which currently elides user tags
    KNOWN_FAILURE;
  }
  else
  {
    thrust::device_vector<int> vec(9001);

    thrust::sequence(vec.begin(), vec.end());
    thrust::reverse(vec.begin(), vec.end());

    // call something we know will invoke get_temporary_buffer
    my_memory_system sys(0);

    thrust::sort(sys, vec.begin(), vec.end());

    ASSERT_EQUAL(true, thrust::is_sorted(vec.begin(), vec.end()));
    ASSERT_EQUAL(true, sys.is_valid());
  }
}
DECLARE_UNITTEST(TestGetTemporaryBufferDispatchImplicit);

#endif

void TestTemporaryBufferOldCustomization()
{
  using system           = my_old_namespace::my_old_temporary_allocation_system;
  using pointer          = thrust::pointer<int, system>;
  using pointer_and_size = cuda::std::pair<pointer, std::ptrdiff_t>;

  system sys;

  {
    pointer_and_size ps = thrust::get_temporary_buffer<int>(sys, 0);

    // The magic values are defined in `my_old_namespace` above.
    ASSERT_EQUAL(ps.first.get(), reinterpret_cast<int*>(4217));
    ASSERT_EQUAL(ps.second, 314);

    thrust::return_temporary_buffer(sys, ps.first, ps.second);
  }
}
DECLARE_UNITTEST(TestTemporaryBufferOldCustomization);

void TestTemporaryBufferNewCustomization()
{
  using system           = my_new_namespace::my_new_temporary_allocation_system;
  using pointer          = thrust::pointer<int, system>;
  using pointer_and_size = cuda::std::pair<pointer, std::ptrdiff_t>;

  system sys;

  {
    pointer_and_size ps = thrust::get_temporary_buffer<int>(sys, 0);

    // The magic values are defined in `my_new_namespace` above.
    ASSERT_EQUAL(ps.first.get(), reinterpret_cast<int*>(1742));
    ASSERT_EQUAL(ps.second, 413);

    thrust::return_temporary_buffer(sys, ps.first, ps.second);
  }
}
DECLARE_UNITTEST(TestTemporaryBufferNewCustomization);