/*
 * Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
 *
 * SPDX-License-Identifier: MIT
 */

#include "memcpy2d_tests_common.hh"

#include <hip_test_common.hh>
#include <hip/hip_runtime_api.h>
#include <resource_guards.hh>
#include <utils.hh>

HIP_TEST_CASE(Unit_hipMemcpyParam2DAsync_Positive_Basic) {
  using namespace std::placeholders;

  constexpr bool async = true;

  const auto stream_type = GENERATE(Streams::nullstream, Streams::perThread, Streams::created);
  const StreamGuard stream_guard(stream_type);
  const hipStream_t stream = stream_guard.stream();

  SECTION("Device to Host") {
    Memcpy2DDeviceToHostShell<async>(
        std::bind(MemcpyParam2DAdapter<async>(), _1, _2, _3, _4, _5, _6, _7, stream), stream);
  }
  SECTION("Device to Device") {
    SECTION("Peer access disabled") {
      Memcpy2DDeviceToDeviceShell<async, false>(
          std::bind(MemcpyParam2DAdapter<async>(), _1, _2, _3, _4, _5, _6, _7, stream), stream);
    }
    SECTION("Peer access enabled") {
      Memcpy2DDeviceToDeviceShell<async, true>(
          std::bind(MemcpyParam2DAdapter<async>(), _1, _2, _3, _4, _5, _6, _7, stream), stream);
    }
  }
  SECTION("Host to Device") {
    Memcpy2DHostToDeviceShell<async>(
        std::bind(MemcpyParam2DAdapter<async>(), _1, _2, _3, _4, _5, _6, _7, stream), stream);
  }
  SECTION("Host to Host") {
    Memcpy2DHostToHostShell<async>(
        std::bind(MemcpyParam2DAdapter<async>(), _1, _2, _3, _4, _5, _6, _7, stream), stream);
  }
}

HIP_TEST_CASE(Unit_hipMemcpyParam2DAsync_Positive_Synchronization_Behavior) {
  using namespace std::placeholders;

  constexpr bool async = true;

  HIP_CHECK(hipDeviceSynchronize());

  SECTION("Host to Device") {
    Memcpy2DHtoDSyncBehavior(
        std::bind(MemcpyParam2DAdapter<async>(), _1, _2, _3, _4, _5, _6, _7, nullptr), false);
  }
#if HT_NVIDIA  // Disabled on AMD due to defect - EXSWHTEC-233
  SECTION("Device to Pageable Host") {
    Memcpy2DDtoHPageableSyncBehavior(
        std::bind(MemcpyParam2DAdapter<async>(), _1, _2, _3, _4, _5, _6, _7, nullptr), true);
  }
#endif
#if HT_NVIDIA  // Disabled on AMD due to defect - EXSWHTEC-236
  SECTION("Device to Pinned Host") {
    Memcpy2DDtoHPinnedSyncBehavior(
        std::bind(MemcpyParam2DAdapter<async>(), _1, _2, _3, _4, _5, _6, _7, nullptr), false);
  }
#endif
  SECTION("Device to Device") {
    Memcpy2DDtoDSyncBehavior(
        std::bind(MemcpyParam2DAdapter<async>(), _1, _2, _3, _4, _5, _6, _7, nullptr), false);
  }
#if HT_NVIDIA  // Disabled on AMD due to defect - EXSWHTEC-233
  SECTION("Host to Host") {
    Memcpy2DHtoHSyncBehavior(
        std::bind(MemcpyParam2DAdapter<async>(), _1, _2, _3, _4, _5, _6, _7, nullptr), true);
  }
#endif
}

HIP_TEST_CASE(Unit_hipMemcpyParam2DAsync_Positive_Parameters) {
  constexpr bool async = true;
  Memcpy2DZeroWidthHeight<async>(MemcpyParam2DAdapter<async>());
}

HIP_TEST_CASE(Unit_hipMemcpyParam2DAsync_Positive_Array) {
  CHECK_IMAGE_SUPPORT
  constexpr bool async = true;
  SECTION("Array from/to Host") {
    MemcpyParam2DArrayHostShell<async>(MemcpyParam2DAdapter<async>());
  }
  SECTION("Array from/to Device") {
    MemcpyParam2DArrayDeviceShell<async>(MemcpyParam2DAdapter<async>());
  }
}

HIP_TEST_CASE(Unit_hipMemcpyParam2DAsync_Negative_Parameters) {
  constexpr bool async = true;

  constexpr size_t cols = 128;
  constexpr size_t rows = 128;

  constexpr auto NegativeTests = [](void* dst, size_t dpitch, void* src, size_t spitch,
                                    size_t width, size_t height, hipMemcpyKind kind) {
    SECTION("dst == nullptr") {
      HIP_CHECK_ERROR(MemcpyParam2DAdapter<async>()(static_cast<void*>(nullptr), dpitch, src,
                                                    spitch, width, height, kind),
                      hipErrorInvalidValue);
    }
    SECTION("src == nullptr") {
      HIP_CHECK_ERROR(MemcpyParam2DAdapter<async>()(dst, dpitch, static_cast<void*>(nullptr),
                                                    spitch, width, height, kind),
                      hipErrorInvalidValue);
    }
    SECTION("dstPitch < WidthInBytes") {
      HIP_CHECK_ERROR(
          MemcpyParam2DAdapter<async>()(dst, width - 1, src, spitch, width, height, kind),
          hipErrorInvalidValue);
    }
    SECTION("srcPitch < WidthInBytes") {
      HIP_CHECK_ERROR(
          MemcpyParam2DAdapter<async>()(dst, dpitch, src, width - 1, width, height, kind),
          hipErrorInvalidValue);
    }
    SECTION("dpitch > max pitch") {
      int attr = 0;
      HIP_CHECK(hipDeviceGetAttribute(&attr, hipDeviceAttributeMaxPitch, 0));
      HIP_CHECK_ERROR(MemcpyParam2DAdapter<async>()(dst, static_cast<size_t>(attr) + 1, src, spitch,
                                                    width, height, kind),
                      hipErrorInvalidValue);
    }
    SECTION("spitch > max pitch") {
      int attr = 0;
      HIP_CHECK(hipDeviceGetAttribute(&attr, hipDeviceAttributeMaxPitch, 0));
      HIP_CHECK_ERROR(MemcpyParam2DAdapter<async>()(dst, dpitch, src, static_cast<size_t>(attr) + 1,
                                                    width, height, kind),
                      hipErrorInvalidValue);
    }
#if HT_NVIDIA  // Disabled on AMD due to defect - EXSWHTEC-237
    SECTION("WidthInBytes + srcXInBytes > srcPitch") {
      HIP_CHECK_ERROR(MemcpyParam2DAdapter<async>(make_hipExtent(spitch - width + 1, 0, 0))(
                          dst, dpitch, src, spitch, width, height, kind),
                      hipErrorInvalidValue);
    }
    SECTION("WidthInBytes + dstXInBytes > dstPitch") {
      HIP_CHECK_ERROR(MemcpyParam2DAdapter<async>(make_hipExtent(0, 0, 0),
                                                  make_hipExtent(dpitch - width + 1, 0, 0))(
                          dst, dpitch, src, spitch, width, height, kind),
                      hipErrorInvalidValue);
    }
    SECTION("srcY out of bounds") {
      HIP_CHECK_ERROR(MemcpyParam2DAdapter<async>(make_hipExtent(0, 1, 0))(dst, dpitch, src, spitch,
                                                                           width, height, kind),
                      hipErrorInvalidValue);
    }
    SECTION("dstY out of bounds") {
      HIP_CHECK_ERROR(MemcpyParam2DAdapter<async>(make_hipExtent(0, 0, 0), make_hipExtent(0, 1, 0))(
                          dst, dpitch, src, spitch, width, height, kind),
                      hipErrorInvalidValue);
    }
#endif
  };

  SECTION("Host to device") {
    LinearAllocGuard2D<int> device_alloc(cols, rows);
    LinearAllocGuard<int> host_alloc(LinearAllocs::hipHostMalloc, device_alloc.pitch() * rows);
    NegativeTests(device_alloc.ptr(), device_alloc.pitch(), host_alloc.ptr(), device_alloc.pitch(),
                  device_alloc.width(), device_alloc.height(), hipMemcpyHostToDevice);
  }

  SECTION("Device to host") {
    LinearAllocGuard2D<int> device_alloc(cols, rows);
    LinearAllocGuard<int> host_alloc(LinearAllocs::hipHostMalloc, device_alloc.pitch() * rows);
    NegativeTests(host_alloc.ptr(), device_alloc.pitch(), device_alloc.ptr(), device_alloc.pitch(),
                  device_alloc.width(), device_alloc.height(), hipMemcpyDeviceToHost);
  }

  SECTION("Host to host") {
    LinearAllocGuard<int> src_alloc(LinearAllocs::hipHostMalloc, cols * rows * sizeof(int));
    LinearAllocGuard<int> dst_alloc(LinearAllocs::hipHostMalloc, cols * rows * sizeof(int));
    NegativeTests(dst_alloc.ptr(), cols * sizeof(int), src_alloc.ptr(), cols * sizeof(int),
                  cols * sizeof(int), rows, hipMemcpyHostToHost);
  }

  SECTION("Device to device") {
    LinearAllocGuard2D<int> src_alloc(cols, rows);
    LinearAllocGuard2D<int> dst_alloc(cols, rows);
    NegativeTests(dst_alloc.ptr(), dst_alloc.pitch(), src_alloc.ptr(), src_alloc.pitch(),
                  dst_alloc.width(), dst_alloc.height(), hipMemcpyDeviceToDevice);
  }
}

static constexpr size_t NUM_W{10};
static constexpr size_t NUM_H{10};

HIP_TEST_CASE(Unit_hipMemcpyParam2DAsync_Capture) {
  void* device_a = nullptr;
  void* device_b = nullptr;
  size_t pitch_a = 0;
  size_t pitch_b = 0;
  constexpr size_t kWidthInBytes = NUM_W * sizeof(int);
  constexpr size_t kHeight = NUM_H;

  HIP_CHECK(hipMallocPitch(&device_a, &pitch_a, kWidthInBytes, kHeight));
  HIP_CHECK(hipMallocPitch(&device_b, &pitch_b, kWidthInBytes, kHeight));

  hipStream_t stream = nullptr;
  HIP_CHECK(hipStreamCreate(&stream));

  hip_Memcpy2D memcpy_desc{};
  memcpy_desc.srcMemoryType = hipMemoryTypeDevice;
  memcpy_desc.srcDevice = reinterpret_cast<hipDeviceptr_t>(device_a);
  memcpy_desc.srcPitch = pitch_a;
  memcpy_desc.dstMemoryType = hipMemoryTypeDevice;
  memcpy_desc.dstDevice = reinterpret_cast<hipDeviceptr_t>(device_b);
  memcpy_desc.dstPitch = pitch_b;
  memcpy_desc.WidthInBytes = kWidthInBytes;
  memcpy_desc.Height = kHeight;

  GENERATE_CAPTURE();
  BEGIN_CAPTURE(stream);
  HIP_CHECK(hipMemcpyParam2DAsync(&memcpy_desc, stream));
  END_CAPTURE(stream);

  HIP_CHECK(hipStreamDestroy(stream));
  HIP_CHECK(hipFree(device_a));
  HIP_CHECK(hipFree(device_b));
}