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

#include <hip_test_common.hh>
#include <vector>

const size_t width = 1024;
const size_t height = 1;
const int N = width * height;
const int value = 120;

/**
 * Test Description
 * ------------------------
 * - Test case to validate basic functionality of
 *   hipDrvGraphExecMemsetNodeSetParams with below scenario,
 *   1) Allocate device memory, host memory
 *   2) Create Graph
 *   3) Add Memset node to Graph to set the elements value to 120
 *   4) Add Memcpy node to Graph to copy from device to host memory
 *   4) Take GraphExec and Instantiate the graph
 *   5) Launch the Graph using GraphExec
 *   6) Validate the host array and all elements should contain value 120
 *   7) Prapare new Memset node parameters with value to 121
 *   8) Update the Node parameter in the memset node with new Memset node
 *      parameters using hipDrvGraphExecMemsetNodeSetParams
 *   9) Launch the Graph using GraphExec
 *   10) Again Validate the host array and all elements should contain
 *       value 121
 * Test source
 * ------------------------
 * - catch/unit/memory/hipDrvGraphExecMemsetNodeSetParams.cc
 */
HIP_TEST_CASE(Unit_hipDrvGraphExecMemsetNodeSetParams_BasicPositive) {
  HIP_CHECK(hipInit(0));
  hipDevice_t device;
  hipCtx_t context;
  HIP_CHECK(hipDeviceGet(&device, 0));
  HIP_CHECK(hipCtxCreate(&context, 0, device));

  size_t pitch;
  hipDeviceptr_t devMemSrc;
  HIP_CHECK(hipMallocPitch(reinterpret_cast<void**>(&devMemSrc), &pitch, width, height));

  char* hostMemDst = new char[N];
  REQUIRE(hostMemDst != nullptr);
  for (int i = 0; i < N; i++) {
    hostMemDst[i] = 0;
  }

  hipGraph_t graph = nullptr;
  HIP_CHECK(hipGraphCreate(&graph, 0));

  hipGraphNode_t memsetNode, memcpyNode;

  // Prepare memset node
  hipMemsetParams initialMemsetParams{};
  initialMemsetParams.dst = reinterpret_cast<void*>(devMemSrc);
  initialMemsetParams.pitch = pitch;
  initialMemsetParams.elementSize = sizeof(char);
  initialMemsetParams.width = width;
  initialMemsetParams.height = height;
  initialMemsetParams.value = value;

  HIP_CHECK(
      hipDrvGraphAddMemsetNode(&memsetNode, graph, nullptr, 0, &initialMemsetParams, context));

  // Prepare memcpyNode
  ::std::vector<hipGraphNode_t> memcpyNodeDependencies;
  memcpyNodeDependencies.push_back(memsetNode);

  HIP_MEMCPY3D memcpyParams{};
  memcpyParams.srcMemoryType = hipMemoryTypeDevice;
  memcpyParams.srcDevice = devMemSrc;
  memcpyParams.srcPitch = pitch;
  memcpyParams.dstMemoryType = hipMemoryTypeHost;
  memcpyParams.dstHost = hostMemDst;
  memcpyParams.dstPitch = width;
  memcpyParams.srcXInBytes = 0;
  memcpyParams.srcY = 0;
  memcpyParams.srcZ = 0;
  memcpyParams.dstXInBytes = 0;
  memcpyParams.dstY = 0;
  memcpyParams.dstZ = 0;
  memcpyParams.WidthInBytes = width;
  memcpyParams.Height = height;
  memcpyParams.Depth = 1;

  HIP_CHECK(hipDrvGraphAddMemcpyNode(&memcpyNode, graph, memcpyNodeDependencies.data(),
                                     memcpyNodeDependencies.size(), &memcpyParams, context));

  hipGraphExec_t graphExec;
  HIP_CHECK(hipGraphInstantiate(&graphExec, graph, nullptr, nullptr, 0));
  HIP_CHECK(hipGraphLaunch(graphExec, nullptr));
  HIP_CHECK(hipStreamSynchronize(0));

  for (int i = 0; i < N; i++) {
    REQUIRE(hostMemDst[i] == value);
  }

  hipMemsetParams newMemsetParams{};
  newMemsetParams.dst = reinterpret_cast<void*>(devMemSrc);
  newMemsetParams.pitch = pitch;
  newMemsetParams.elementSize = sizeof(char);
  newMemsetParams.width = width;
  newMemsetParams.height = height;
  newMemsetParams.value = value + 1;

  // Update with new memset node params
  HIP_CHECK(hipDrvGraphExecMemsetNodeSetParams(graphExec, memsetNode, &newMemsetParams, context));

  for (int i = 0; i < N; i++) {
    hostMemDst[i] = 0;
  }

  HIP_CHECK(hipGraphLaunch(graphExec, nullptr));
  HIP_CHECK(hipStreamSynchronize(0));

  for (int i = 0; i < N; i++) {
    REQUIRE(hostMemDst[i] == (value + 1));
  }

  HIP_CHECK(hipGraphExecDestroy(graphExec));
  HIP_CHECK(hipGraphDestroy(graph));
  HIP_CHECK(hipFree(reinterpret_cast<void*>(devMemSrc)));
  delete[] hostMemDst;
  HIP_CHECK(hipCtxDestroy(context));
}

/**
 * Test Description
 * ------------------------
 * - Test case to validate below scenarios,
 *    1) With Invalid Graph Exec
 *    2) With Invalid Node
 *    3) With Invalid Node params
 *    4) With Invalid destination address in Node params
 * ------------------------
 * - catch/unit/memory/hipDrvGraphExecMemsetNodeSetParams.cc
 */
HIP_TEST_CASE(Unit_hipDrvGraphExecMemsetNodeSetParams_Negative) {
  HIP_CHECK(hipInit(0));
  hipDevice_t device;
  hipCtx_t context;
  HIP_CHECK(hipDeviceGet(&device, 0));
  HIP_CHECK(hipCtxCreate(&context, 0, device));

  hipGraph_t graph = nullptr;
  HIP_CHECK(hipGraphCreate(&graph, 0));

  size_t pitch;
  hipDeviceptr_t devMemSrc;
  HIP_CHECK(hipMallocPitch(reinterpret_cast<void**>(&devMemSrc), &pitch, width, height));

  // Prepare memset node
  hipMemsetParams memsetParams{};
  memsetParams.dst = reinterpret_cast<void*>(devMemSrc);
  memsetParams.pitch = pitch;
  memsetParams.elementSize = sizeof(char);
  memsetParams.width = width;
  memsetParams.height = height;
  memsetParams.value = value;

  hipGraphNode_t memsetNode;
  HIP_CHECK(hipDrvGraphAddMemsetNode(&memsetNode, graph, nullptr, 0, &memsetParams, context));

  hipGraphExec_t graphExec;
  HIP_CHECK(hipGraphInstantiate(&graphExec, graph, nullptr, nullptr, 0));

  SECTION("Invalid Graph Exec") {
    HIP_CHECK_ERROR(hipDrvGraphExecMemsetNodeSetParams(nullptr, memsetNode, &memsetParams, context),
                    hipErrorInvalidValue);
  }
  SECTION("Invalid Node") {
    HIP_CHECK_ERROR(hipDrvGraphExecMemsetNodeSetParams(graphExec, nullptr, &memsetParams, context),
                    hipErrorInvalidValue);
  }
  SECTION("Invalid Node params") {
    HIP_CHECK_ERROR(hipDrvGraphExecMemsetNodeSetParams(graphExec, memsetNode, nullptr, context),
                    hipErrorInvalidValue);
  }
  SECTION("Invalid destination address in Node params") {
    hipMemsetParams memsetParams{};
    memsetParams.dst = 0;
    memsetParams.pitch = pitch;
    memsetParams.elementSize = sizeof(char);
    memsetParams.width = width;
    memsetParams.height = height;
    memsetParams.value = value;

    HIP_CHECK_ERROR(hipDrvGraphExecMemsetNodeSetParams(graphExec, memsetNode, nullptr, context),
                    hipErrorInvalidValue);
  }

  HIP_CHECK(hipGraphExecDestroy(graphExec));
  HIP_CHECK(hipGraphDestroy(graph));
  HIP_CHECK(hipFree(reinterpret_cast<void*>(devMemSrc)));
  HIP_CHECK(hipCtxDestroy(context));
}