/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *  * Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 *  * Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *  * Neither the name of NVIDIA CORPORATION nor the names of its
 *    contributors may be used to endorse or promote products derived
 *    from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
 * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
 * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
 * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include "cudaNvSci.h"
#include <cuda.h>
#include <condition_variable>
#include <iostream>
#include <thread>

std::mutex m_mutex;
std::condition_variable m_condVar;
bool workSubmitted = false;

class cudaNvSciSignal {
 private:
  NvSciSyncModule m_syncModule;
  NvSciBufModule m_bufModule;

  NvSciSyncAttrList m_syncAttrList;
  NvSciSyncFence *m_fence;

  NvSciBufAttrList m_rawBufAttrList;
  NvSciBufAttrList m_imageBufAttrList;
  NvSciBufAttrList m_buffAttrListOut[2];
  NvSciBufAttrKeyValuePair pairArrayOut[10];

  cudaExternalMemory_t extMemRawBuf, extMemImageBuf;
  cudaMipmappedArray_t d_mipmapArray;
  cudaArray_t d_mipLevelArray;
  cudaTextureObject_t texObject;
  cudaExternalSemaphore_t signalSem;

  cudaStream_t streamToRun;
  int m_cudaDeviceId;
  CUuuid m_devUUID;
  uint64_t m_imageWidth;
  uint64_t m_imageHeight;
  void *d_outputBuf;
  size_t m_bufSize;

 public:
  cudaNvSciSignal(NvSciBufModule bufModule, NvSciSyncModule syncModule,
                  int cudaDeviceId, int bufSize, uint64_t imageWidth,
                  uint64_t imageHeight, NvSciSyncFence *fence)
      : m_syncModule(syncModule),
        m_bufModule(bufModule),
        m_cudaDeviceId(cudaDeviceId),
        d_outputBuf(NULL),
        m_bufSize(bufSize),
        m_imageWidth(imageWidth),
        m_imageHeight(imageHeight),
        m_fence(fence) {
    initCuda();

    checkNvSciErrors(NvSciSyncAttrListCreate(m_syncModule, &m_syncAttrList));
    checkNvSciErrors(NvSciBufAttrListCreate(m_bufModule, &m_rawBufAttrList));
    checkNvSciErrors(NvSciBufAttrListCreate(m_bufModule, &m_imageBufAttrList));

    setRawBufAttrList(m_bufSize);
    setImageBufAttrList(m_imageWidth, m_imageHeight);

    checkCudaErrors(cudaDeviceGetNvSciSyncAttributes(
        m_syncAttrList, m_cudaDeviceId, cudaNvSciSyncAttrSignal));
  }

  ~cudaNvSciSignal() {
    checkCudaErrors(cudaSetDevice(m_cudaDeviceId));
    checkCudaErrors(cudaFreeMipmappedArray(d_mipmapArray));
    checkCudaErrors(cudaFree(d_outputBuf));
    checkCudaErrors(cudaDestroyExternalSemaphore(signalSem));
    checkCudaErrors(cudaDestroyExternalMemory(extMemRawBuf));
    checkCudaErrors(cudaDestroyExternalMemory(extMemImageBuf));
    checkCudaErrors(cudaDestroyTextureObject(texObject));
    checkCudaErrors(cudaStreamDestroy(streamToRun));
  }

  void initCuda() {
    checkCudaErrors(cudaSetDevice(m_cudaDeviceId));
    checkCudaErrors(
        cudaStreamCreateWithFlags(&streamToRun, cudaStreamNonBlocking));

    int major = 0, minor = 0;
    checkCudaErrors(cudaDeviceGetAttribute(
        &major, cudaDevAttrComputeCapabilityMajor, m_cudaDeviceId));
    checkCudaErrors(cudaDeviceGetAttribute(
        &minor, cudaDevAttrComputeCapabilityMinor, m_cudaDeviceId));
    printf(
        "[cudaNvSciSignal] GPU Device %d: \"%s\" with compute capability "
        "%d.%d\n\n",
        m_cudaDeviceId, _ConvertSMVer2ArchName(major, minor), major, minor);

#ifdef cuDeviceGetUuid_v2
    CUresult res = cuDeviceGetUuid_v2(&m_devUUID, m_cudaDeviceId);
#else
    CUresult res = cuDeviceGetUuid(&m_devUUID, m_cudaDeviceId);
#endif

    if (res != CUDA_SUCCESS) {
      fprintf(stderr, "Driver API error = %04d \n", res);
      exit(EXIT_FAILURE);
    }
  }

  void setRawBufAttrList(uint64_t size) {
    NvSciBufType bufType = NvSciBufType_RawBuffer;
    bool cpuAccess = false;
    NvSciBufAttrValAccessPerm perm = NvSciBufAccessPerm_ReadWrite;
    NvSciBufAttrKeyValuePair rawBufAttrs[] = {
        {NvSciBufRawBufferAttrKey_Size, &size, sizeof(size)},
        {NvSciBufGeneralAttrKey_Types, &bufType, sizeof(bufType)},
        {NvSciBufGeneralAttrKey_NeedCpuAccess, &cpuAccess, sizeof(cpuAccess)},
        {NvSciBufGeneralAttrKey_RequiredPerm, &perm, sizeof(perm)},
        {NvSciBufGeneralAttrKey_GpuId, &m_devUUID, sizeof(m_devUUID)},
    };

    checkNvSciErrors(NvSciBufAttrListSetAttrs(
        m_rawBufAttrList, rawBufAttrs,
        sizeof(rawBufAttrs) / sizeof(NvSciBufAttrKeyValuePair)));
  }

  void setImageBufAttrList(uint32_t width, uint32_t height) {
    NvSciBufType bufType = NvSciBufType_Image;
    NvSciBufAttrValImageLayoutType layout = NvSciBufImage_BlockLinearType;
    NvSciBufAttrValAccessPerm perm = NvSciBufAccessPerm_ReadWrite;

    uint32_t planeCount = 1;
    uint32_t planeWidths[] = {width};
    uint32_t planeHeights[] = {height};
    uint64_t lrpad = 0, tbpad = 100;

    bool cpuAccessFlag = false;

    NvSciBufAttrValColorFmt planecolorfmts[] = {NvSciColor_B8G8R8A8};
    NvSciBufAttrValColorStd planecolorstds[] = {NvSciColorStd_SRGB};
    NvSciBufAttrValImageScanType planescantype[] = {NvSciBufScan_InterlaceType};

    NvSciBufAttrKeyValuePair imgBufAttrs[] = {
        {NvSciBufGeneralAttrKey_Types, &bufType, sizeof(bufType)},
        {NvSciBufImageAttrKey_PlaneCount, &planeCount, sizeof(planeCount)},
        {NvSciBufImageAttrKey_Layout, &layout, sizeof(layout)},
        {NvSciBufImageAttrKey_TopPadding, &tbpad, sizeof(tbpad)},
        {NvSciBufImageAttrKey_BottomPadding, &tbpad, sizeof(tbpad)},
        {NvSciBufImageAttrKey_LeftPadding, &lrpad, sizeof(lrpad)},
        {NvSciBufImageAttrKey_RightPadding, &lrpad, sizeof(lrpad)},
        {NvSciBufImageAttrKey_PlaneColorFormat, planecolorfmts,
         sizeof(planecolorfmts)},
        {NvSciBufImageAttrKey_PlaneColorStd, planecolorstds,
         sizeof(planecolorstds)},
        {NvSciBufImageAttrKey_PlaneWidth, planeWidths, sizeof(planeWidths)},
        {NvSciBufImageAttrKey_PlaneHeight, planeHeights, sizeof(planeHeights)},
        {NvSciBufGeneralAttrKey_NeedCpuAccess, &cpuAccessFlag,
         sizeof(cpuAccessFlag)},
        {NvSciBufGeneralAttrKey_RequiredPerm, &perm, sizeof(perm)},
        {NvSciBufImageAttrKey_PlaneScanType, planescantype,
         sizeof(planescantype)},
        {NvSciBufGeneralAttrKey_GpuId, &m_devUUID, sizeof(m_devUUID)},
    };

    checkNvSciErrors(NvSciBufAttrListSetAttrs(
        m_imageBufAttrList, imgBufAttrs,
        sizeof(imgBufAttrs) / sizeof(NvSciBufAttrKeyValuePair)));
  }

  NvSciSyncAttrList getNvSciSyncAttrList() { return m_syncAttrList; }

  NvSciBufAttrList getNvSciRawBufAttrList() { return m_rawBufAttrList; }

  NvSciBufAttrList getNvSciImageBufAttrList() { return m_imageBufAttrList; }

  void runRotateImageAndSignal(unsigned char *imageData) {
    int numOfGPUs = 0;
    checkCudaErrors(cudaGetDeviceCount(&numOfGPUs));  // For cuda init purpose
    checkCudaErrors(cudaSetDevice(m_cudaDeviceId));

    copyDataToImageArray(imageData);
    createTexture();

    float angle = 0.5f;  // angle to rotate image by (in radians)
    rotateKernel(texObject, angle, (unsigned int *)d_outputBuf, m_imageWidth,
                 m_imageHeight, streamToRun);

    signalExternalSemaphore();
  }

  void cudaImportNvSciSemaphore(NvSciSyncObj syncObj) {
    checkCudaErrors(cudaSetDevice(m_cudaDeviceId));

    cudaExternalSemaphoreHandleDesc extSemDesc;
    memset(&extSemDesc, 0, sizeof(extSemDesc));
    extSemDesc.type = cudaExternalSemaphoreHandleTypeNvSciSync;
    extSemDesc.handle.nvSciSyncObj = (void *)syncObj;

    checkCudaErrors(cudaImportExternalSemaphore(&signalSem, &extSemDesc));
  }

  void signalExternalSemaphore() {
    cudaExternalSemaphoreSignalParams signalParams;
    memset(&signalParams, 0, sizeof(signalParams));
    // For cross-process signaler-waiter applications need to use NvSciIpc
    // and NvSciSync[Export|Import] utilities to share the NvSciSyncFence
    // across process. This step is optional in single-process.
    signalParams.params.nvSciSync.fence = (void *)m_fence;
    signalParams.flags = 0;

    checkCudaErrors(cudaSignalExternalSemaphoresAsync(&signalSem, &signalParams,
                                                      1, streamToRun));
  }

  void cudaImportNvSciRawBuf(NvSciBufObj inputBufObj) {
    checkCudaErrors(cudaSetDevice(m_cudaDeviceId));
    checkNvSciErrors(
        NvSciBufObjGetAttrList(inputBufObj, &m_buffAttrListOut[0]));

    memset(pairArrayOut, 0, sizeof(NvSciBufAttrKeyValuePair) * 10);
    pairArrayOut[0].key = NvSciBufRawBufferAttrKey_Size;

    checkNvSciErrors(
        NvSciBufAttrListGetAttrs(m_buffAttrListOut[0], pairArrayOut, 1));

    uint64_t size = *(uint64_t *)pairArrayOut[0].value;

    cudaExternalMemoryHandleDesc memHandleDesc;
    memset(&memHandleDesc, 0, sizeof(memHandleDesc));
    memHandleDesc.type = cudaExternalMemoryHandleTypeNvSciBuf;
    memHandleDesc.handle.nvSciBufObject = inputBufObj;
    memHandleDesc.size = size;
    checkCudaErrors(cudaImportExternalMemory(&extMemRawBuf, &memHandleDesc));

    cudaExternalMemoryBufferDesc bufferDesc;
    memset(&bufferDesc, 0, sizeof(bufferDesc));
    bufferDesc.offset = 0;
    bufferDesc.size = size;
    m_bufSize = size;
    checkCudaErrors(cudaExternalMemoryGetMappedBuffer(
        &d_outputBuf, extMemRawBuf, &bufferDesc));
  }

  void cudaImportNvSciImage(NvSciBufObj inputBufObj) {
    checkCudaErrors(cudaSetDevice(m_cudaDeviceId));
    checkNvSciErrors(
        NvSciBufObjGetAttrList(inputBufObj, &m_buffAttrListOut[1]));

    memset(pairArrayOut, 0, sizeof(NvSciBufAttrKeyValuePair) * 10);
    pairArrayOut[0].key = NvSciBufImageAttrKey_Size;
    pairArrayOut[1].key = NvSciBufImageAttrKey_Alignment;
    pairArrayOut[2].key = NvSciBufImageAttrKey_PlaneCount;
    pairArrayOut[3].key = NvSciBufImageAttrKey_PlaneWidth;
    pairArrayOut[4].key = NvSciBufImageAttrKey_PlaneHeight;

    checkNvSciErrors(
        NvSciBufAttrListGetAttrs(m_buffAttrListOut[1], pairArrayOut, 5));

    uint64_t size = *(uint64_t *)pairArrayOut[0].value;
    uint64_t alignment = *(uint64_t *)pairArrayOut[1].value;
    uint64_t planeCount = *(uint64_t *)pairArrayOut[2].value;
    uint64_t imageWidth = *(uint64_t *)pairArrayOut[3].value;
    uint64_t imageHeight = *(uint64_t *)pairArrayOut[4].value;

    cudaExternalMemoryHandleDesc memHandleDesc;
    memset(&memHandleDesc, 0, sizeof(memHandleDesc));
    memHandleDesc.type = cudaExternalMemoryHandleTypeNvSciBuf;
    memHandleDesc.handle.nvSciBufObject = inputBufObj;
    memHandleDesc.size = size;
    checkCudaErrors(cudaImportExternalMemory(&extMemImageBuf, &memHandleDesc));

    cudaExtent extent = {};
    memset(&extent, 0, sizeof(extent));
    extent.width = imageWidth;
    extent.height = imageHeight;
    extent.depth = 0;

    cudaChannelFormatDesc desc;
    desc.x = 8;
    desc.y = 8;
    desc.z = 8;
    desc.w = 8;
    desc.f = cudaChannelFormatKindUnsigned;

    cudaExternalMemoryMipmappedArrayDesc mipmapDesc = {0};
    mipmapDesc.offset = 0;
    mipmapDesc.formatDesc = desc;
    mipmapDesc.extent = extent;
    mipmapDesc.flags = 0;

    mipmapDesc.numLevels = 1;
    checkCudaErrors(cudaExternalMemoryGetMappedMipmappedArray(
        &d_mipmapArray, extMemImageBuf, &mipmapDesc));
  }

  void copyDataToImageArray(unsigned char *imageData) {
    uint32_t mipLevelId = 0;
    checkCudaErrors(cudaGetMipmappedArrayLevel(&d_mipLevelArray, d_mipmapArray,
                                               mipLevelId));

    checkCudaErrors(cudaMemcpy2DToArrayAsync(
        d_mipLevelArray, 0, 0, imageData, m_imageWidth * sizeof(unsigned int),
        m_imageWidth * sizeof(unsigned int), m_imageHeight,
        cudaMemcpyHostToDevice, streamToRun));
  }

  void createTexture() {
    cudaResourceDesc texRes;
    memset(&texRes, 0, sizeof(cudaResourceDesc));

    texRes.resType = cudaResourceTypeArray;
    texRes.res.array.array = d_mipLevelArray;

    cudaTextureDesc texDescr;
    memset(&texDescr, 0, sizeof(cudaTextureDesc));

    texDescr.normalizedCoords = true;
    texDescr.filterMode = cudaFilterModeLinear;
    texDescr.addressMode[0] = cudaAddressModeWrap;
    texDescr.addressMode[1] = cudaAddressModeWrap;
    texDescr.readMode = cudaReadModeNormalizedFloat;

    checkCudaErrors(
        cudaCreateTextureObject(&texObject, &texRes, &texDescr, NULL));
  }
};

class cudaNvSciWait {
 private:
  NvSciSyncModule m_syncModule;
  NvSciBufModule m_bufModule;

  NvSciSyncAttrList m_syncAttrList;
  NvSciBufAttrList m_rawBufAttrList;
  NvSciBufAttrList m_buffAttrListOut;
  NvSciBufAttrKeyValuePair pairArrayOut[10];
  NvSciSyncFence *m_fence;

  cudaExternalMemory_t extMemRawBuf;
  cudaExternalSemaphore_t waitSem;
  cudaStream_t streamToRun;
  int m_cudaDeviceId;
  CUuuid m_devUUID;
  void *d_outputBuf;
  size_t m_bufSize;
  size_t imageWidth;
  size_t imageHeight;

 public:
  cudaNvSciWait(NvSciBufModule bufModule, NvSciSyncModule syncModule,
                int cudaDeviceId, int bufSize, NvSciSyncFence *fence)
      : m_bufModule(bufModule),
        m_syncModule(syncModule),
        m_cudaDeviceId(cudaDeviceId),
        m_bufSize(bufSize),
        m_fence(fence) {
    initCuda();
    checkNvSciErrors(NvSciSyncAttrListCreate(m_syncModule, &m_syncAttrList));
    checkNvSciErrors(NvSciBufAttrListCreate(m_bufModule, &m_rawBufAttrList));

    setRawBufAttrList(m_bufSize);
    checkCudaErrors(cudaDeviceGetNvSciSyncAttributes(
        m_syncAttrList, m_cudaDeviceId, cudaNvSciSyncAttrWait));
  }

  ~cudaNvSciWait() {
    checkCudaErrors(cudaStreamDestroy(streamToRun));
    checkCudaErrors(cudaDestroyExternalSemaphore(waitSem));
    checkCudaErrors(cudaDestroyExternalMemory(extMemRawBuf));
    checkCudaErrors(cudaFree(d_outputBuf));
  }

  void initCuda() {
    checkCudaErrors(cudaSetDevice(m_cudaDeviceId));
    checkCudaErrors(
        cudaStreamCreateWithFlags(&streamToRun, cudaStreamNonBlocking));
#ifdef cuDeviceGetUuid_v2
    CUresult res = cuDeviceGetUuid_v2(&m_devUUID, m_cudaDeviceId);
#else
    CUresult res = cuDeviceGetUuid(&m_devUUID, m_cudaDeviceId);
#endif
    if (res != CUDA_SUCCESS) {
      fprintf(stderr, "Driver API error = %04d \n", res);
      exit(EXIT_FAILURE);
    }

    int major = 0, minor = 0;
    checkCudaErrors(cudaDeviceGetAttribute(
        &major, cudaDevAttrComputeCapabilityMajor, m_cudaDeviceId));
    checkCudaErrors(cudaDeviceGetAttribute(
        &minor, cudaDevAttrComputeCapabilityMinor, m_cudaDeviceId));
    printf(
        "[cudaNvSciWait] GPU Device %d: \"%s\" with compute capability "
        "%d.%d\n\n",
        m_cudaDeviceId, _ConvertSMVer2ArchName(major, minor), major, minor);
  }

  void setRawBufAttrList(uint64_t size) {
    NvSciBufType bufType = NvSciBufType_RawBuffer;
    bool cpuAccess = false;
    NvSciBufAttrValAccessPerm perm = NvSciBufAccessPerm_ReadWrite;
    NvSciBufAttrKeyValuePair rawBufAttrs[] = {
        {NvSciBufRawBufferAttrKey_Size, &size, sizeof(size)},
        {NvSciBufGeneralAttrKey_Types, &bufType, sizeof(bufType)},
        {NvSciBufGeneralAttrKey_NeedCpuAccess, &cpuAccess, sizeof(cpuAccess)},
        {NvSciBufGeneralAttrKey_RequiredPerm, &perm, sizeof(perm)},
        {NvSciBufGeneralAttrKey_GpuId, &m_devUUID, sizeof(m_devUUID)},
    };

    checkNvSciErrors(NvSciBufAttrListSetAttrs(
        m_rawBufAttrList, rawBufAttrs,
        sizeof(rawBufAttrs) / sizeof(NvSciBufAttrKeyValuePair)));
  }

  NvSciSyncAttrList getNvSciSyncAttrList() { return m_syncAttrList; }

  NvSciBufAttrList getNvSciRawBufAttrList() { return m_rawBufAttrList; }

  void runImageGrayscale(std::string image_filename, size_t imageWidth,
                         size_t imageHeight) {
    int numOfGPUs = 0;
    checkCudaErrors(cudaGetDeviceCount(&numOfGPUs));  // For cuda init purpose
    checkCudaErrors(cudaSetDevice(m_cudaDeviceId));

    waitExternalSemaphore();
    launchGrayScaleKernel((unsigned int *)d_outputBuf, image_filename,
                          imageWidth, imageHeight, streamToRun);
  }

  void cudaImportNvSciSemaphore(NvSciSyncObj syncObj) {
    checkCudaErrors(cudaSetDevice(m_cudaDeviceId));

    cudaExternalSemaphoreHandleDesc extSemDesc;
    memset(&extSemDesc, 0, sizeof(extSemDesc));
    extSemDesc.type = cudaExternalSemaphoreHandleTypeNvSciSync;
    extSemDesc.handle.nvSciSyncObj = (void *)syncObj;

    checkCudaErrors(cudaImportExternalSemaphore(&waitSem, &extSemDesc));
  }

  void waitExternalSemaphore() {
    checkCudaErrors(cudaSetDevice(m_cudaDeviceId));

    cudaExternalSemaphoreWaitParams waitParams;
    memset(&waitParams, 0, sizeof(waitParams));
    // For cross-process signaler-waiter applications need to use NvSciIpc
    // and NvSciSync[Export|Import] utilities to share the NvSciSyncFence
    // across process. This step is optional in single-process.
    waitParams.params.nvSciSync.fence = (void *)m_fence;
    waitParams.flags = 0;

    checkCudaErrors(
        cudaWaitExternalSemaphoresAsync(&waitSem, &waitParams, 1, streamToRun));
  }

  void cudaImportNvSciRawBuf(NvSciBufObj inputBufObj) {
    checkCudaErrors(cudaSetDevice(m_cudaDeviceId));

    checkNvSciErrors(NvSciBufObjGetAttrList(inputBufObj, &m_buffAttrListOut));

    memset(pairArrayOut, 0, sizeof(NvSciBufAttrKeyValuePair) * 10);
    pairArrayOut[0].key = NvSciBufRawBufferAttrKey_Size;

    checkNvSciErrors(
        NvSciBufAttrListGetAttrs(m_buffAttrListOut, pairArrayOut, 1));

    uint64_t size = *(uint64_t *)pairArrayOut[0].value;

    cudaExternalMemoryHandleDesc memHandleDesc;
    memset(&memHandleDesc, 0, sizeof(memHandleDesc));
    memHandleDesc.type = cudaExternalMemoryHandleTypeNvSciBuf;
    memHandleDesc.handle.nvSciBufObject = inputBufObj;
    memHandleDesc.size = size;
    checkCudaErrors(cudaImportExternalMemory(&extMemRawBuf, &memHandleDesc));

    cudaExternalMemoryBufferDesc bufferDesc;
    memset(&bufferDesc, 0, sizeof(bufferDesc));
    bufferDesc.offset = 0;
    bufferDesc.size = size;
    m_bufSize = size;

    checkCudaErrors(cudaExternalMemoryGetMappedBuffer(
        &d_outputBuf, extMemRawBuf, &bufferDesc));
  }
};

void thread_rotateAndSignal(cudaNvSciSignal *cudaNvSciSignalObj,
                            unsigned char *imageData) {
  std::lock_guard<std::mutex> guard(m_mutex);
  cudaNvSciSignalObj->runRotateImageAndSignal(imageData);
  workSubmitted = true;
  m_condVar.notify_one();
}

void thread_waitAndGrayscale(cudaNvSciWait *cudaNvSciWaitObj,
                             std::string image_filename, size_t imageWidth,
                             size_t imageHeight) {
  // Acquire the lock
  std::unique_lock<std::mutex> mlock(m_mutex);
  m_condVar.wait(mlock, [] { return workSubmitted; });
  cudaNvSciWaitObj->runImageGrayscale(image_filename, imageWidth, imageHeight);
}

cudaNvSci::cudaNvSci(int isMultiGPU, std::vector<int> &deviceIds,
                     unsigned char *imageData, size_t width, size_t height)
    : m_isMultiGPU(isMultiGPU),
      image_data(imageData),
      imageWidth(width),
      imageHeight(height) {
  if (isMultiGPU) {
    m_cudaNvSciSignalDeviceId = deviceIds[0];
    m_cudaNvSciWaitDeviceId = deviceIds[1];
  } else {
    m_cudaNvSciSignalDeviceId = m_cudaNvSciWaitDeviceId = deviceIds[0];
  }

  m_bufSize = imageWidth * imageHeight * sizeof(unsigned int);
}

void cudaNvSci::initNvSci() {
  checkNvSciErrors(NvSciSyncModuleOpen(&syncModule));
  checkNvSciErrors(NvSciBufModuleOpen(&buffModule));
  fence = (NvSciSyncFence *)calloc(1, sizeof(NvSciSyncFence));
}

void cudaNvSci::runCudaNvSci(std::string &image_filename) {
  initNvSci();

  cudaNvSciSignal rotateAndSignal(buffModule, syncModule,
                                  m_cudaNvSciSignalDeviceId, m_bufSize,
                                  imageWidth, imageHeight, fence);
  cudaNvSciWait waitAndGrayscale(buffModule, syncModule,
                                 m_cudaNvSciWaitDeviceId, m_bufSize, fence);

  rawBufUnreconciledList[0] = rotateAndSignal.getNvSciRawBufAttrList();
  rawBufUnreconciledList[1] = waitAndGrayscale.getNvSciRawBufAttrList();

  createNvSciRawBufObj();

  imageBufUnreconciledList[0] = rotateAndSignal.getNvSciImageBufAttrList();

  createNvSciBufImageObj();

  rotateAndSignal.cudaImportNvSciRawBuf(rawBufObj);
  rotateAndSignal.cudaImportNvSciImage(imageBufObj);

  waitAndGrayscale.cudaImportNvSciRawBuf(rawBufObj);

  syncUnreconciledList[0] = rotateAndSignal.getNvSciSyncAttrList();
  syncUnreconciledList[1] = waitAndGrayscale.getNvSciSyncAttrList();

  createNvSciSyncObj();

  rotateAndSignal.cudaImportNvSciSemaphore(syncObj);
  waitAndGrayscale.cudaImportNvSciSemaphore(syncObj);

  std::thread rotateThread(&thread_rotateAndSignal, &rotateAndSignal,
                           image_data);

  std::thread grayscaleThread(&thread_waitAndGrayscale, &waitAndGrayscale,
                              image_filename, imageWidth, imageHeight);

  rotateThread.join();
  grayscaleThread.join();
}

void cudaNvSci::createNvSciRawBufObj() {
  int numAttrList = 2;
  checkNvSciErrors(NvSciBufAttrListReconcile(rawBufUnreconciledList,
                                             numAttrList, &rawBufReconciledList,
                                             &buffConflictList));
  checkNvSciErrors(NvSciBufObjAlloc(rawBufReconciledList, &rawBufObj));
  printf("created NvSciBufObj\n");
}

void cudaNvSci::createNvSciBufImageObj() {
  int numAttrList = 1;
  checkNvSciErrors(NvSciBufAttrListReconcile(
      imageBufUnreconciledList, numAttrList, &imageBufReconciledList,
      &imageBufConflictList));
  checkNvSciErrors(NvSciBufObjAlloc(imageBufReconciledList, &imageBufObj));
  printf("created NvSciBufImageObj\n");
}

void cudaNvSci::createNvSciSyncObj() {
  int numAttrList = 2;
  checkNvSciErrors(NvSciSyncAttrListReconcile(syncUnreconciledList, numAttrList,
                                              &syncReconciledList,
                                              &syncConflictList));
  checkNvSciErrors(NvSciSyncObjAlloc(syncReconciledList, &syncObj));
  printf("created NvSciSyncObj\n");
}
