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

/**
Testcase Scenarios :
 1) Validates parameter list of hipStreamAddCallback.
 2) Validates hipStreamAddCallback functionality with default stream.
 3) Validates hipStreamAddCallback functionality with defined stream.
*/

#include <hip_test_common.hh>
#include <hip_test_kernels.hh>
#include <chrono>
#include <thread>

#define UNUSED(expr)                                                                               \
  do {                                                                                             \
    (void)(expr);                                                                                  \
  } while (0)

#ifdef __HIP_PLATFORM_AMD__
#define HIPRT_CB
#endif

namespace hipStreaAddCallbackTest {
constexpr size_t NSize = 4 * 1024 * 1024;
float *A_h, *C_h;
bool gcbDone = false;
bool gPassed = true;
void* ptr0xff = reinterpret_cast<void*>(0xffffffff);
void* gusrptr;
hipStream_t gstream;

void HIPRT_CB Callback(hipStream_t stream, hipError_t status, void* userData) {
  UNUSED(stream);
  HIP_CHECK(status);
  REQUIRE(userData == NULL);
  gPassed = true;
  for (size_t i = 0; i < NSize; i++) {
    if (C_h[i] != A_h[i] * A_h[i]) {
      UNSCOPED_INFO("Data mismatch :" << i);
      gPassed = false;
      break;
    }
  }
  gcbDone = true;
}
/**
 * Validates functionality of hipStreamAddCallback with default/created stream.
 */
bool testStreamCallbackFunctionality(bool isDefault) {
  float *A_d, *C_d;
  size_t Nbytes = NSize * sizeof(float);
  gcbDone = false;
  A_h = reinterpret_cast<float*>(malloc(Nbytes));
  REQUIRE(A_h != nullptr);
  C_h = reinterpret_cast<float*>(malloc(Nbytes));
  REQUIRE(C_h != nullptr);

  // Fill with Phi + i
  for (size_t i = 0; i < NSize; i++) {
    A_h[i] = 1.618f + i;
  }

  HIP_CHECK(hipMalloc(&A_d, Nbytes));
  HIP_CHECK(hipMalloc(&C_d, Nbytes));
  if (isDefault) {
    HIP_CHECK(hipMemcpyAsync(A_d, A_h, Nbytes, hipMemcpyHostToDevice, 0));

    const unsigned threadsPerBlock = 1024;
    const int blocks = (NSize % threadsPerBlock == 0) ? (NSize / threadsPerBlock)
                                                      : ((NSize / threadsPerBlock) + 1);
    hipLaunchKernelGGL((HipTest::vector_square), dim3(blocks), dim3(threadsPerBlock), 0, 0, A_d,
                       C_d, NSize);
    HIP_CHECK(hipGetLastError());
    HIP_CHECK(hipMemcpyAsync(C_h, C_d, Nbytes, hipMemcpyDeviceToHost, 0));
    HIP_CHECK(hipStreamAddCallback(0, Callback, nullptr, 0));
    while (!gcbDone)
      std::this_thread::sleep_for(std::chrono::microseconds(100000));  // Sleep for 100 ms
  } else {
    hipStream_t mystream;
    HIP_CHECK(hipStreamCreateWithFlags(&mystream, hipStreamNonBlocking));

    HIP_CHECK(hipMemcpyAsync(A_d, A_h, Nbytes, hipMemcpyHostToDevice, mystream));

    const unsigned threadsPerBlock = 1024;
    const int blocks = (NSize % threadsPerBlock == 0) ? (NSize / threadsPerBlock)
                                                      : ((NSize / threadsPerBlock) + 1);
    hipLaunchKernelGGL((HipTest::vector_square), dim3(blocks), dim3(threadsPerBlock), 0, mystream,
                       A_d, C_d, NSize);
    HIP_CHECK(hipGetLastError());
    HIP_CHECK(hipMemcpyAsync(C_h, C_d, Nbytes, hipMemcpyDeviceToHost, mystream));
    HIP_CHECK(hipStreamAddCallback(mystream, Callback, nullptr, 0));
    while (!gcbDone)
      std::this_thread::sleep_for(std::chrono::microseconds(100000));  // Sleep for 100 ms
    HIP_CHECK(hipStreamDestroy(mystream));
  }
  HIP_CHECK(hipFree(reinterpret_cast<void*>(C_d)));
  HIP_CHECK(hipFree(reinterpret_cast<void*>(A_d)));
  free(C_h);
  free(A_h);
  return gPassed;
}
/**
 * Scenario1: Validates if callback = nullptr returns error code for created stream.
 * Scenario2: Validates if callback = nullptr returns error code for default stream.
 * Scenario3: Validates if flag != 0 returns error code for created stream.
 * Scenario4: Validates if flag != 0 returns error code for default stream.
 * Scenario5: Validates if userData pointer is passed properly to callback.
 * Scenario6: Validates if stream value is passed properly to callback.
 */
void Callback_ChkUsrdataPtr(hipStream_t stream, hipError_t status, void* userData) {
  REQUIRE(stream == gstream);
  HIP_CHECK(status);
  gPassed = true;
  if (gusrptr != userData) {
    gPassed = false;
  }
  gcbDone = true;
}

void Callback_ChkStreamValue(hipStream_t stream, hipError_t status, void* userData) {
  REQUIRE(userData == nullptr);
  HIP_CHECK(status);
  gPassed = true;
  if (stream != gstream) {
    gPassed = false;
  }
  gcbDone = true;
}
}  // namespace hipStreaAddCallbackTest


using hipStreaAddCallbackTest::Callback;
using hipStreaAddCallbackTest::Callback_ChkStreamValue;
using hipStreaAddCallbackTest::Callback_ChkUsrdataPtr;
using hipStreaAddCallbackTest::gcbDone;
using hipStreaAddCallbackTest::gPassed;
using hipStreaAddCallbackTest::gstream;
using hipStreaAddCallbackTest::gusrptr;
using hipStreaAddCallbackTest::ptr0xff;
using hipStreaAddCallbackTest::testStreamCallbackFunctionality;


/*
 * Validates parameter list of hipStreamAddCallback.
 */
HIP_TEST_CASE(Unit_hipStreamAddCallback_ParamTst_Positive) {
  hipStream_t mystream;
  HIP_CHECK(hipStreamCreate(&mystream));

  // Scenario1
  SECTION("userData pointer value validation") {
    gstream = mystream;
    gusrptr = ptr0xff;
    gPassed = true;
    gcbDone = false;
    HIP_CHECK(hipStreamAddCallback(mystream, Callback_ChkUsrdataPtr, gusrptr, 0));
    while (!gcbDone) {
      std::this_thread::sleep_for(std::chrono::microseconds(100000));  // Sleep for 100 ms
    }
    REQUIRE(gPassed);
  }
  // Scenario2
  SECTION("stream value validation") {
    gstream = mystream;
    gPassed = true;
    gcbDone = false;
    HIP_CHECK(hipStreamAddCallback(mystream, Callback_ChkStreamValue, nullptr, 0));
    while (!gcbDone) {
      std::this_thread::sleep_for(std::chrono::microseconds(100000));  // Sleep for 100 ms
    }
    REQUIRE(gPassed);
  }
  HIP_CHECK(hipStreamDestroy(mystream));
}

/*
 * Negative tests for validation of hipStreamAddCallback parameter list.
 */
HIP_TEST_CASE(Unit_hipStreamAddCallback_ParamTst_Negative) {
  hipStream_t mystream;
  HIP_CHECK(hipStreamCreate(&mystream));

  // Scenario1
  SECTION("callback is nullptr for non-default stream") {
    REQUIRE_FALSE(hipSuccess == hipStreamAddCallback(mystream, nullptr, nullptr, 0));
  }
  // Scenario2
  SECTION("callback is nullptr for default stream") {
    REQUIRE_FALSE(hipSuccess == hipStreamAddCallback(0, nullptr, nullptr, 0));
  }
  // Scenario3
  SECTION("flag is nonzero for non-default stream") {
    REQUIRE_FALSE(hipSuccess == hipStreamAddCallback(mystream, Callback, nullptr, 10));
  }
  // Scenario4
  SECTION("flag is nonzero for default stream") {
    REQUIRE_FALSE(hipSuccess == hipStreamAddCallback(0, Callback, nullptr, 10));
  }
  HIP_CHECK(hipStreamDestroy(mystream));
}

/*
 * Validates hipStreamAddCallback functionality with default stream.
 */
HIP_TEST_CASE(Unit_hipStreamAddCallback_WithDefaultStream) {
  bool TestPassed = true;
  TestPassed = testStreamCallbackFunctionality(true);
  REQUIRE(TestPassed);
}

/*
 * Validates hipStreamAddCallback functionality with defined stream.
 */
HIP_TEST_CASE(Unit_hipStreamAddCallback_WithCreatedStream) {
  bool TestPassed = true;
  TestPassed = testStreamCallbackFunctionality(false);
  REQUIRE(TestPassed);
}