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

#include <hip_test_common.hh>
#include <hip_test_checkers.hh>

#ifdef __linux__
#include <sys/types.h>
#include <sys/mman.h>
#include <sys/wait.h>
#include <fcntl.h>
#include <semaphore.h>
#include <unistd.h>

/**
 * @addtogroup hipIpcOpenMemHandle hipIpcOpenMemHandle
 * @{
 * @ingroup DeviceTest
 * `hipIpcOpenMemHandle(void** devPtr, hipIpcMemHandle_t handle, unsigned int flags)` -
 * Opens an interprocess memory handle exported from another process
 * and returns a device pointer usable in the local process.
 */

#define NUM_ELMTS 1024
#define NUM_THREADS 10


typedef struct mem_handle {
  int device;
  hipIpcMemHandle_t memHandle;
  bool IfTestPassed;
} hip_ipc_t;


// This testcase verifies the hipIpcMemAccess APIs as follows
// The following program spawns a child process and does the following
// Parent iterate through each device, create memory -- create hipIpcMemhandle
// stores the mem handle in mmaped memory, release the child using sem_post()
// and wait for child to release itself(parent process)
// child process:
// Child process get the ipc mem handle using hipIpcOpenMemHandle
// Iterate through all the available gpus and do Device to Device copies
// and check for data consistencies and close the hipIpcCloseMemHandle
// release the parent and wait for parent to release itself(child)

/**
 * Test Description
 * ------------------------
 *  - Verifies that getting and opening mem handle works correctly
 *    in specific scenarion, and handles the case when the same device
 *    is used in both processes.
 *  - Creates memory from the parent process for each device.
 *  - Spawns child process and waits for it to finish.
 *  - Child process gets the handle and check data consistencies.
 * Test source
 * ------------------------
 *  - unit/multiproc/hipIpcMemAccessTest.cc
 * Test requirements
 * ------------------------
 *  - Host specific (LINUX)
 *  - HIP_VERSION >= 5.2
 */
HIP_TEST_CASE(Unit_hipIpcMemAccess_Semaphores) {
  hip_ipc_t* shrd_mem = NULL;
  pid_t pid;
  size_t N = 1024;
  size_t Nbytes = N * sizeof(int);
  int *A_d{nullptr}, *B_d{nullptr}, *C_d{nullptr};
  int *A_h{nullptr}, *C_h{nullptr};
  sem_t *sem_ob1{nullptr}, *sem_ob2{nullptr};
  int Num_devices = 0, CanAccessPeer = 0;

  std::string cmd_line = "rm -rf /dev/shm/sem.my-sem-object*";
  int res = system(cmd_line.c_str());
  REQUIRE(res != -1);
  sem_ob1 = sem_open("/my-sem-object1", O_CREAT | O_EXCL, 0660, 0);
  sem_ob2 = sem_open("/my-sem-object2", O_CREAT | O_EXCL, 0660, 0);
  REQUIRE(sem_ob1 != SEM_FAILED);
  REQUIRE(sem_ob2 != SEM_FAILED);

  shrd_mem = reinterpret_cast<hip_ipc_t*>(
      mmap(NULL, sizeof(hip_ipc_t), PROT_READ | PROT_WRITE, MAP_SHARED | MAP_ANONYMOUS, 0, 0));
  REQUIRE(shrd_mem != NULL);
  shrd_mem->IfTestPassed = true;
  HipTest::initArrays<int>(nullptr, nullptr, nullptr, &A_h, nullptr, &C_h, N, false);
  pid = fork();
  if (pid != 0) {
    // Parent process
    HIP_CHECK(hipGetDeviceCount(&Num_devices));
    for (int i = 0; i < Num_devices; ++i) {
      if (shrd_mem->IfTestPassed == true) {
        HIP_CHECK(hipSetDevice(i));
        HIP_CHECK(hipMalloc(&A_d, Nbytes));
        HIP_CHECK(
            hipIpcGetMemHandle(reinterpret_cast<hipIpcMemHandle_t*>(&shrd_mem->memHandle), A_d));
        HIP_CHECK(hipMemcpy(A_d, A_h, Nbytes, hipMemcpyHostToDevice));
        shrd_mem->device = i;
        if ((sem_post(sem_ob1)) == -1) {
          // Need to use inline function to release resources.
          shrd_mem->IfTestPassed = false;
          WARN("sem_post() call failed in parent process.");
        }
        if ((sem_wait(sem_ob2)) == -1) {
          shrd_mem->IfTestPassed = false;
          WARN("sem_wait() call failed in parent process.");
        }
        HIP_CHECK(hipFree(A_d));
      }
    }
  } else {
    // Child process
    HIP_CHECK(hipGetDeviceCount(&Num_devices));
    for (int j = 0; j < Num_devices; ++j) {
      HIP_CHECK(hipSetDevice(j));
      if ((sem_wait(sem_ob1)) == -1) {
        shrd_mem->IfTestPassed = false;
        WARN("sem_wait() call failed in child process.");
        if ((sem_post(sem_ob2)) == -1) {
          shrd_mem->IfTestPassed = false;
          WARN("sem_post() call on sem_ob2 failed");
          exit(1);
        }
      }
      for (int i = 0; i < Num_devices; ++i) {

        HIP_CHECK(hipDeviceCanAccessPeer(&CanAccessPeer, i, shrd_mem->device));
        if (CanAccessPeer == 1) {
          HIP_CHECK(hipDeviceEnablePeerAccess(i, 0));
          HIP_CHECK(hipSetDevice(i));
          HIP_CHECK(hipMalloc(&C_d, Nbytes));
          HIP_CHECK(hipIpcOpenMemHandle(reinterpret_cast<void**>(&B_d), shrd_mem->memHandle,
                                        hipIpcMemLazyEnablePeerAccess));
          HIP_CHECK(hipMemcpy(C_d, B_d, Nbytes, hipMemcpyDeviceToDevice));
          HIP_CHECK(hipMemcpy(C_h, C_d, Nbytes, hipMemcpyDeviceToHost));
          HipTest::checkTest<int>(A_h, C_h, N);
          memset(reinterpret_cast<void*>(C_h), 0, Nbytes);
          // Checking if the data obtained from Ipc shared memory is consistent
          HIP_CHECK(hipMemcpy(C_h, B_d, Nbytes, hipMemcpyDeviceToHost));
          HipTest::checkTest<int>(A_h, C_h, N);
          HIP_CHECK(hipIpcCloseMemHandle(reinterpret_cast<void*>(B_d)));
          HIP_CHECK(hipFree(C_d));
        }
      }
      if ((sem_post(sem_ob2)) == -1) {
        shrd_mem->IfTestPassed = false;
        WARN("sem_post() call on sem_ob2 failed");
        exit(1);
      }
    }
    exit(0);
  }
  if ((sem_unlink("/my-sem-object1")) == -1) {
    WARN("sem_unlink() call on /my-sem-object1 failed");
  }
  if ((sem_unlink("/my-sem-object2")) == -1) {
    WARN("sem_unlink() call on /my-sem-object2 failed");
  }
  int rFlag = 0;
  waitpid(pid, &rFlag, 0);
  REQUIRE(shrd_mem->IfTestPassed == true);
  HipTest::freeArrays<int>(nullptr, nullptr, nullptr, A_h, nullptr, C_h, false);
}

/**
 * Test Description
 * ------------------------
 *  - Validates handling of valid and invalid arguments for
 *    [hipIpcGetMemHandle](@ref hipIpcGetMemHandle):
 *    -# When memory handle pointer is `nullptr`
 *      - Expected output: return `hipErrorInvalidValue`
 *    -# When device pointer is `nullptr`
 *      - Expected output: return `hipErrorInvalidValue`
 *    -# When both pointers are `nullptr`
 *      - Expected output: return `hipErrorInvalidValue`
 *    -# When both pointers are valid
 *      - Expected output: return `hipSuccess`
 *  - Validates handling of valid and invalid arguments for
 *    [hipIpcOpenMemHandle](@ref hipIpcOpenMemHandle):
 *    -# When device pointer is `nullptr`
 *      - Expected output: return `hipErrorInvalidValue`
 *    -# When memory handle pointer uninitialized
 *      - Expected output: return `hipErrorInvalidValue` or `hipErrorInvalidDevicePointer`
 *    -# When memory handle has random flags
 *      - Expected output: return `hipErrorInvalidValue`
 *  - Validates handling of valid and invalid arguments for
 *    [hipIpcCloseMemHandle](@ref hipIpcCloseMemHandle):
 *    -# When device pointer is `nullptr`
 *      - Expected output: return `hipErrorInvalidValue`
 * Test source
 * ------------------------
 *  - unit/multiproc/hipIpcMemAccessTest.cc
 * Test requirements
 * ------------------------
 *  - Host specific (LINUX)
 *  - HIP_VERSION >= 5.2
 */
HIP_TEST_CASE(Unit_hipIpcMemAccess_ParameterValidation) {
  hipIpcMemHandle_t MemHandle;
  hipIpcMemHandle_t MemHandleUninit;
  void *Ad{}, *Ad2{};
  hipError_t ret;

  HIP_CHECK(hipMalloc(&Ad, 1024));

#if HT_AMD
  // Test is disabled for nvidia as api resulting in seg fault.
  SECTION("Get mem handle with handle as nullptr") {
    ret = hipIpcGetMemHandle(nullptr, Ad);
    REQUIRE(ret == hipErrorInvalidValue);
  }
#endif
  SECTION("Get mem handle with devptr as nullptr") {
    ret = hipIpcGetMemHandle(&MemHandle, nullptr);
    REQUIRE(ret == hipErrorInvalidValue);
  }

  SECTION("Get mem handle with handle/devptr as nullptr") {
    ret = hipIpcGetMemHandle(nullptr, nullptr);
    REQUIRE(ret == hipErrorInvalidValue);
  }

  SECTION("Get mem handle with valid devptr") {
    ret = hipIpcGetMemHandle(&MemHandle, Ad);
    REQUIRE(ret == hipSuccess);
  }

  SECTION("Open mem handle with devptr as nullptr") {
    ret = hipIpcOpenMemHandle(nullptr, MemHandle, hipIpcMemLazyEnablePeerAccess);
    REQUIRE(ret == hipErrorInvalidValue);
  }

  SECTION("Open mem handle with handle as un-initialized") {
    ret = hipIpcOpenMemHandle(&Ad2, MemHandleUninit, hipIpcMemLazyEnablePeerAccess);
    REQUIRE((ret == hipErrorInvalidValue || ret == hipErrorInvalidDevicePointer));
  }
#if HT_AMD
  // Test is disabled for nvidia as api not returning expected value.
  SECTION("Open mem handle with flags as random value") {
    constexpr unsigned int flags = 123;
    HIP_CHECK(hipIpcGetMemHandle(&MemHandle, Ad));
    ret = hipIpcOpenMemHandle(&Ad2, MemHandle, flags);
    REQUIRE(ret == hipErrorInvalidValue);
  }
#endif
  SECTION("Close mem handle with devptr(nullptr)") {
    ret = hipIpcCloseMemHandle(nullptr);
    REQUIRE(ret == hipErrorInvalidValue);
  }

  HIP_CHECK(hipFree(Ad));
}

/**
 * End doxygen group hipIpcOpenMemHandle.
 * @}
 */

#endif