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

/*
hipMallocMipmappedArray API test scenarios
1. Basic Functionality
2. Negative Scenarios
3. Allocating Small and big chunk data
4. Multithreaded scenario
*/

#include <array>
#include <hip_test_common.hh>
#include "hipArrayCommon.hh"

static constexpr auto ARRAY_SIZE{4};
static constexpr auto BIG_ARRAY_SIZE{100};
static constexpr auto ARRAY_LOOP{100};


/*
 * This API verifies  memory allocations for small and
 * bigger chunks of data.
 * Two scenarios are verified in this API
 * 1. SmallArray: Allocates ARRAY_SIZE in a loop and
 *    releases the memory and verifies the meminfo.
 * 2. BigArray: Allocates BIG_ARRAY_SIZE in a loop and
 *    releases the memory and verifies the meminfo
 *
 * In both cases, the memory info before allocation and
 * after releasing the memory should be the same
 *
 */
static void MallocMipmappedArray_DiffSizes(int gpu) {
  HIP_CHECK_THREAD(hipSetDevice(gpu));
  // Use of GENERATE in thead function causes random failures with multithread condition.
  std::vector<size_t> runs{ARRAY_SIZE, BIG_ARRAY_SIZE};
  for (const auto& size : runs) {
    auto numLevelsLimit = floor(log2(size));
    for (unsigned int numLevels = 0; numLevels < numLevelsLimit; numLevels++) {
      size_t width{size}, height{size}, depth{size};
      hipChannelFormatDesc channelDesc = hipCreateChannelDesc<float>();
      std::array<hipMipmappedArray_t, ARRAY_LOOP> arr;
      size_t pavail, avail, total;
      HIP_CHECK_THREAD(hipMemGetInfo(&pavail, &total));

      for (int i = 0; i < ARRAY_LOOP; i++) {
        HIP_CHECK_THREAD(hipMallocMipmappedArray(&arr[i], &channelDesc,
                                                 make_hipExtent(width, height, depth),
                                                 (1 + numLevels), hipArrayDefault));
      }
      for (int i = 0; i < ARRAY_LOOP; i++) {
        HIP_CHECK_THREAD(hipFreeMipmappedArray(arr[i]));
      }

      HIP_CHECK_THREAD(hipMemGetInfo(&avail, &total));
      REQUIRE_THREAD(pavail == avail);
    }
  }
}

HIP_TEST_CASE(Unit_hipMallocMipmappedArray_DiffSizes) {
  MallocMipmappedArray_DiffSizes(0);
  HIP_CHECK_THREAD_FINALIZE();
}

/*
This testcase verifies the hipMallocMipmappedArray API in multithreaded
scenario by launching threads in parallel on multiple GPUs
and verifies the hipMallocMipmappedArray API with small and big chunks data
*/
HIP_TEST_CASE(Unit_hipMallocMipmappedArray_MultiThread) {
  std::vector<std::thread> threadlist;
  int devCnt = 0;
  devCnt = HipTest::getDeviceCount();
  for (int i = 0; i < devCnt; i++) {
    threadlist.push_back(std::thread(MallocMipmappedArray_DiffSizes, i));
  }

  for (auto& t : threadlist) {
    t.join();
  }
  HIP_CHECK_THREAD_FINALIZE();
}

namespace {
void checkMipmappedArrayIsExpected(hipArray_t level_array,
                                   const hipChannelFormatDesc& expected_desc,
                                   const hipExtent& expected_extent,
                                   const unsigned int expected_flags) {
// hipArrayGetInfo doesn't currently exist (EXSWCPHIPT-87)
#if HT_AMD
  std::ignore = level_array;
  std::ignore = expected_desc;
  std::ignore = expected_extent;
  std::ignore = expected_flags;
#else
  cudaChannelFormatDesc queried_desc;
  cudaExtent queried_extent;
  unsigned int queried_flags;

  cudaArrayGetInfo(&queried_desc, &queried_extent, &queried_flags, level_array);

  REQUIRE(expected_desc.x == queried_desc.x);
  REQUIRE(expected_desc.y == queried_desc.y);
  REQUIRE(expected_desc.z == queried_desc.z);
  REQUIRE(expected_desc.f == queried_desc.f);

  REQUIRE(expected_extent.width == queried_extent.width);
  REQUIRE(expected_extent.height == queried_extent.height);
  REQUIRE(expected_extent.depth == queried_extent.depth);

  REQUIRE(expected_flags == queried_flags);
#endif
}
}  // namespace

HIP_TEMPLATE_TEST_CASE(Unit_hipMallocMipmappedArray_happy, char, uint2, int4, short4, float) {
  hipMipmappedArray_t array;
  const auto desc = hipCreateChannelDesc<TestType>();
#if HT_AMD
  const unsigned int flags = hipArrayDefault;
#else
  const unsigned int flags =
      GENERATE(hipArrayDefault, hipArraySurfaceLoadStore, hipArrayTextureGather);
#endif
  constexpr size_t size = 64;
  const unsigned int numLevels = GENERATE(1, 3, 5, 7);

  std::vector<hipExtent> extents;
  extents.reserve(3);
  extents.push_back({size, size, 0});  // 2D array
  if (flags != hipArrayTextureGather) {
    extents.push_back({size, 0, 0});        // 1D array
    extents.push_back({size, size, size});  // 3D array
  };

  for (const auto extent : extents) {
    CAPTURE(flags, extent.width, extent.height, extent.depth);
    HIP_CHECK_IGNORED_RETURN(hipMallocMipmappedArray(&array, &desc, extent, numLevels, flags),
                             hipErrorNotSupported);
    hipArray_t hipArray = nullptr;
    HIP_CHECK(hipGetMipmappedArrayLevel(&hipArray, array, 0));
    checkMipmappedArrayIsExpected(hipArray, desc, extent, flags);
    HIP_CHECK(hipFreeMipmappedArray(array));
  }
}

#if HT_AMD
constexpr std::array<unsigned int, 1> validFlags{hipArrayDefault};
#else
constexpr std::array<unsigned int, 9> validFlags{
    hipArrayDefault,
    hipArrayDefault | hipArraySurfaceLoadStore,
    hipArrayLayered,
    hipArrayLayered | hipArraySurfaceLoadStore,
    hipArrayCubemap,
    hipArrayCubemap | hipArrayLayered,
    hipArrayCubemap | hipArraySurfaceLoadStore,
    hipArrayCubemap | hipArrayLayered | hipArraySurfaceLoadStore,
    hipArrayTextureGather};
#endif

hipExtent makeMipmappedExtent(unsigned int flag, size_t s) {
  if (flag == hipArrayTextureGather) {
    return make_hipExtent(s, s, 0);
  }
  return make_hipExtent(s, s, s);
}

// Providing the array pointer as nullptr should return an error
HIP_TEST_CASE(Unit_hipMallocMipmappedArray_Negative_NullArrayPtr) {
  hipChannelFormatDesc desc = hipCreateChannelDesc<float4>();
  unsigned int numLevels = 1;
  constexpr size_t s = 6;

  const auto flag = GENERATE(from_range(std::begin(validFlags), std::end(validFlags)));
  HIP_CHECK_ERROR(
      hipMallocMipmappedArray(nullptr, &desc, makeMipmappedExtent(flag, s), numLevels, flag),
      hipErrorInvalidValue);
}

// Providing the description pointer as nullptr should return an error
HIP_TEST_CASE(Unit_hipMallocMipmappedArray_Negative_NullDescPtr) {
  constexpr size_t s = 6;  // 6 to keep cubemap happy
  unsigned int numLevels = 1;
  hipMipmappedArray_t array;

  const auto flag = GENERATE(from_range(std::begin(validFlags), std::end(validFlags)));

  HIP_CHECK_ERROR(
      hipMallocMipmappedArray(&array, nullptr, makeMipmappedExtent(flag, s), numLevels, flag),
      hipErrorInvalidValue);
}

// Zero width arrays are not allowed
HIP_TEST_CASE(Unit_hipMallocMipmappedArray_Negative_ZeroWidth) {
  constexpr size_t s = 6;  // 6 to keep cubemap happy
  unsigned int numLevels = 1;
  hipMipmappedArray_t array;
  hipChannelFormatDesc desc = hipCreateChannelDesc<float4>();

  const auto flag = GENERATE(from_range(std::begin(validFlags), std::end(validFlags)));

  HIP_CHECK_ERRORS(hipMallocMipmappedArray(&array, &desc, make_hipExtent(0, s, s), numLevels, flag),
                   hipErrorInvalidValue, hipErrorNotSupported);
}

// Zero height arrays are only allowed for 1D arrays and layered arrays
HIP_TEST_CASE(Unit_hipMallocMipmappedArray_Negative_ZeroHeight) {
  constexpr size_t s = 6;  // 6 to keep cubemap happy
  unsigned int numLevels = 1;
  hipMipmappedArray_t array;
  hipChannelFormatDesc desc = hipCreateChannelDesc<float4>();
  std::array<unsigned int, 2> exceptions{hipArrayLayered,
                                         hipArrayLayered | hipArraySurfaceLoadStore};

  const auto flag = GENERATE(from_range(std::begin(validFlags), std::end(validFlags)));

  if (std::find(std::begin(exceptions), std::end(exceptions), flag) == std::end(exceptions)) {
    // flag is not in list of exceptions
    HIP_CHECK_ERRORS(
        hipMallocMipmappedArray(&array, &desc, make_hipExtent(s, 0, s), numLevels, flag),
        hipErrorInvalidValue, hipErrorNotSupported);
  }
}

HIP_TEST_CASE(Unit_hipMallocMipmappedArray_Negative_InvalidFlags) {
  constexpr size_t s = 6;  // 6 to keep cubemap happy
  unsigned int numLevels = 1;
  hipMipmappedArray_t array;
  hipChannelFormatDesc desc = hipCreateChannelDesc<float4>();

#if HT_AMD
  const unsigned int flag = 0xDEADBEEF;
#else
  const unsigned int flag =
      GENERATE(0xDEADBEEF, hipArrayTextureGather | hipArraySurfaceLoadStore,
               hipArrayTextureGather | hipArrayCubemap,
               hipArrayTextureGather | hipArraySurfaceLoadStore | hipArrayCubemap);
#endif

  CAPTURE(flag);

  REQUIRE(std::find(std::begin(validFlags), std::end(validFlags), flag) == std::end(validFlags));

  HIP_CHECK_ERRORS(
      hipMallocMipmappedArray(&array, &desc, makeMipmappedExtent(flag, s), numLevels, flag),
      hipErrorInvalidValue, hipErrorNotSupported);
}

void testInvalidDescriptionMipmapped(hipChannelFormatDesc desc) {
  constexpr size_t s = 6;  // 6 to keep cubemap happy
  unsigned int numLevels = 1;
  hipMipmappedArray_t array;

#if HT_NVIDIA
  hipError_t expectedError = hipErrorInvalidChannelDescriptor;
#else
  hipError_t expectedError = hipErrorInvalidValue;
#endif

  const auto flag = GENERATE(from_range(std::begin(validFlags), std::end(validFlags)));
  HIP_CHECK_ERRORS(
      hipMallocMipmappedArray(&array, &desc, makeMipmappedExtent(flag, s), numLevels, flag),
      expectedError, hipErrorNotSupported);
}

HIP_TEST_CASE(Unit_hipMallocMipmappedArray_Negative_InvalidFormat) {
  hipChannelFormatDesc desc = hipCreateChannelDesc<float4>();
  desc.f = GENERATE(hipChannelFormatKindNone, 0xBEEF);
  testInvalidDescriptionMipmapped(desc);
}

HIP_TEST_CASE(Unit_hipMallocMipmappedArray_Negative_BadChannelLayout) {
  const int bits = GENERATE(8, 16, 32);
  const hipChannelFormatKind formatKind =
      GENERATE(hipChannelFormatKindSigned, hipChannelFormatKindUnsigned, hipChannelFormatKindFloat);
  if (bits == 8 && formatKind == hipChannelFormatKindFloat) return;


  hipChannelFormatDesc desc = GENERATE_COPY(hipCreateChannelDesc(bits, bits, bits, 0, formatKind),
                                            hipCreateChannelDesc(0, bits, bits, 0, formatKind),
                                            hipCreateChannelDesc(0, bits, bits, bits, formatKind),
                                            hipCreateChannelDesc(bits, 0, bits, 0, formatKind),
                                            hipCreateChannelDesc(bits, bits, 0, bits, formatKind),
                                            hipCreateChannelDesc(0, 0, bits, 0, formatKind),
                                            hipCreateChannelDesc(0, 0, bits, bits, formatKind));

  INFO("kind: " << channelFormatString(formatKind));
  INFO("x: " << desc.x << ", y: " << desc.y << ", z: " << desc.z << ", w: " << desc.w);

  testInvalidDescriptionMipmapped(desc);
}

HIP_TEST_CASE(Unit_hipMallocMipmappedArray_Negative_8BitFloat) {
  hipChannelFormatDesc desc = GENERATE(hipCreateChannelDesc(8, 0, 0, 0, hipChannelFormatKindFloat),
                                       hipCreateChannelDesc(8, 8, 0, 0, hipChannelFormatKindFloat),
                                       hipCreateChannelDesc(8, 8, 8, 8, hipChannelFormatKindFloat));

  testInvalidDescriptionMipmapped(desc);
}

HIP_TEST_CASE(Unit_hipMallocMipmappedArray_Negative_DifferentChannelSizes) {
  const int bitsX = GENERATE(8, 16, 32);
  const int bitsY = GENERATE(8, 16, 32);
  const int bitsZ = GENERATE(8, 16, 32);
  const int bitsW = GENERATE(8, 16, 32);
  if (bitsX == bitsY && bitsY == bitsZ && bitsZ == bitsW) return;  // skip when they are equal

  const hipChannelFormatKind channelFormat =
      GENERATE(hipChannelFormatKindSigned, hipChannelFormatKindUnsigned, hipChannelFormatKindFloat);

  if (channelFormat == hipChannelFormatKindFloat &&
      (bitsX == 8 || bitsY == 8 || bitsZ == 8 || bitsW == 8))
    return;  // 8 bit floats aren't allowed

  hipChannelFormatDesc desc = hipCreateChannelDesc(bitsX, bitsY, bitsZ, bitsW, channelFormat);

  INFO("format: " << channelFormatString(channelFormat) << ", x bits: " << bitsX
                  << ", y bits: " << bitsY << ", z bits: " << bitsZ << ", w bits: " << bitsW);


  testInvalidDescriptionMipmapped(desc);
}

HIP_TEST_CASE(Unit_hipMallocMipmappedArray_Negative_BadChannelSize) {
  const int badBits = GENERATE(-1, 0, 10, 100);
  const hipChannelFormatKind formatKind =
      GENERATE(hipChannelFormatKindSigned, hipChannelFormatKindUnsigned, hipChannelFormatKindFloat);
  hipChannelFormatDesc desc = hipCreateChannelDesc(badBits, badBits, badBits, badBits, formatKind);

  INFO("Number of bits: " << badBits);

  testInvalidDescriptionMipmapped(desc);
}


// hipMallocMipmappedArray should handle the max numeric value gracefully.
HIP_TEST_CASE(Unit_hipMallocMipmappedArray_Negative_NumericLimit) {
  hipMipmappedArray_t arrayPtr;
  unsigned int numLevels = 1;
  hipChannelFormatDesc desc = hipCreateChannelDesc<float>();

  size_t size = std::numeric_limits<size_t>::max();
  const auto flag = GENERATE(from_range(std::begin(validFlags), std::end(validFlags)));
  HIP_CHECK_ERRORS(
      hipMallocMipmappedArray(&arrayPtr, &desc, makeMipmappedExtent(flag, size), numLevels, flag),
      hipErrorInvalidValue, hipErrorNotSupported);
}

// texture gather arrays are only allowed to be 2D
HIP_TEMPLATE_TEST_CASE(Unit_hipMallocMipmappedArray_Negative_Non2DTextureGather, char, uchar2,
                   float2) {
#if HT_AMD
  HIP_SKIP_TEST(HipTest::SkipReason::kTextureGatherUnsupportedAmd);
#endif
  hipMipmappedArray_t array;
  unsigned int numLevels = 1;
  const auto desc = hipCreateChannelDesc<TestType>();

  constexpr unsigned int flags = hipArrayTextureGather;
  constexpr size_t size = 64;
  const hipExtent extent = GENERATE(make_hipExtent(size, 0, 0), make_hipExtent(size, size, size));

  HIP_CHECK_ERROR(hipMallocMipmappedArray(&array, &desc, extent, numLevels, flags),
                  hipErrorInvalidValue);
}

HIP_TEST_CASE(Unit_hipMallocMipmappedArray_Negative_NumLevels) {
  hipMipmappedArray_t array;
  constexpr size_t size = 6;
  unsigned int numLevels = floor(log2(size)) + 2;
  hipChannelFormatDesc desc = hipCreateChannelDesc<float>();

  const auto flag = hipArrayDefault;
#if HT_AMD
  HIP_CHECK_ERRORS(
      hipMallocMipmappedArray(&array, &desc, makeMipmappedExtent(flag, size), numLevels, flag),
      hipErrorInvalidValue, hipErrorNotSupported);
#else
  HIP_CHECK(
      hipMallocMipmappedArray(&array, &desc, makeMipmappedExtent(flag, size), numLevels, flag));
#endif
}

HIP_TEST_CASE(Unit_hipGetMipmappedArrayLevel_Negative) {
  constexpr size_t s = 6;
  unsigned int numLevels = 1;
  hipMipmappedArray_t array;
  hipArray_t level_array;

  hipChannelFormatDesc desc = hipCreateChannelDesc<float>();

  HIP_CHECK_IGNORED_RETURN(
      hipMallocMipmappedArray(&array, &desc, make_hipExtent(s, s, s), numLevels, hipArrayDefault),
      hipErrorNotSupported);
  SECTION("Level is invalid") {
    HIP_CHECK_ERROR(hipGetMipmappedArrayLevel(&level_array, array, 3), hipErrorInvalidValue);
  }
  SECTION("Mipmapped array is nullptr") {
    HIP_CHECK_ERROR(hipGetMipmappedArrayLevel(&level_array, nullptr, 1),
                    hipErrorInvalidResourceHandle);
  }
  HIP_CHECK(hipFreeMipmappedArray(array));
}