/* 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 <string.h>
#include <iostream>
/* Nvidia headers */
#include "nvmedia_utils/cmdline.h"
#include "nvmedia_image.h"
#include "nvmedia_2d.h"
#include "nvmedia_surface.h"
#include "nvmedia_utils/image_utils.h"
#include "nvmedia_image_nvscibuf.h"
#include "nvmedia_producer.h"
#include "nvmedia_2d_nvscisync.h"
#include "nvsci_setup.h"

NvMediaImage *NvMediaImageCreateUsingNvScibuf(NvMediaDevice *device,
                                              NvMediaSurfaceType type,
                                              const NvMediaSurfAllocAttr *attrs,
                                              uint32_t numAttrs, uint32_t flags,
                                              NvSciBufObj &bufobj,
                                              int cudaDeviceId) {
  NvSciBufModule module = NULL;
  NvSciError err = NvSciError_Success;
  NvMediaStatus status = NVMEDIA_STATUS_OK;
  NvSciBufAttrList attrlist = NULL;
  NvSciBufAttrList conflictlist = NULL;
  NvSciBufAttrValAccessPerm access_perm = NvSciBufAccessPerm_ReadWrite;
  NvSciBufAttrKeyValuePair attr_kvp = {NvSciBufGeneralAttrKey_RequiredPerm,
                                       &access_perm, sizeof(access_perm)};
  NvSciBufAttrKeyValuePair pairArrayOut[10];

  NvMediaImage *image = NULL;

  err = NvSciBufModuleOpen(&module);
  if (err != NvSciError_Success) {
    printf("%s: NvSciBuffModuleOpen failed. Error: %d \n", __func__, err);
    goto fail_cleanup;
  }

  err = NvSciBufAttrListCreate(module, &attrlist);
  if (err != NvSciError_Success) {
    printf("%s: SciBufAttrListCreate failed. Error: %d \n", __func__, err);
    goto fail_cleanup;
  }

  err = NvSciBufAttrListSetAttrs(attrlist, &attr_kvp, 1);
  if (err != NvSciError_Success) {
    printf("%s: AccessPermSetAttr failed. Error: %d \n", __func__, err);
    goto fail_cleanup;
  }

  status =
      NvMediaImageFillNvSciBufAttrs(device, type, attrs, numAttrs, 0, attrlist);

  if (status != NVMEDIA_STATUS_OK) {
    printf("%s: ImageFillSciBufAttrs failed. Error: %d \n", __func__, err);
    goto fail_cleanup;
  }

  setupNvSciBuf(bufobj, attrlist, cudaDeviceId);

  status = NvMediaImageCreateFromNvSciBuf(device, bufobj, &image);

  if (status != NVMEDIA_STATUS_OK) {
    printf("%s: ImageCreatefromSciBuf failed. Error: %d \n", __func__, err);
    goto fail_cleanup;
  }

  NvSciBufAttrListFree(attrlist);

  if (module != NULL) {
    NvSciBufModuleClose(module);
  }

  return image;

fail_cleanup:
  if (attrlist != NULL) {
    NvSciBufAttrListFree(attrlist);
  }
  if (bufobj != NULL) {
    NvSciBufObjFree(bufobj);
    bufobj = NULL;
  }

  if (module != NULL) {
    NvSciBufModuleClose(module);
  }
  NvMediaImageDestroy(image);
  return NULL;
}

/* Create NvMediaImage surface based on the input attributes.
 * Returns NVMEDIA_STATUS_OK on success
 */
static NvMediaStatus createSurface(Blit2DTest *ctx,
                                   NvMediaSurfFormatAttr *surfFormatAttrs,
                                   NvMediaSurfAllocAttr *surfAllocAttrs,
                                   uint32_t numSurfAllocAttrs,
                                   NvMediaImage **image, NvSciBufObj &bufObj,
                                   int cudaDeviceId) {
  NvMediaSurfaceType surfType;

  /* create source image */
  surfType =
      NvMediaSurfaceFormatGetType(surfFormatAttrs, NVM_SURF_FMT_ATTR_MAX);
  *image = NvMediaImageCreateUsingNvScibuf(ctx->device, /* device */
                                           surfType,    /* surface type */
                                           surfAllocAttrs, numSurfAllocAttrs, 0,
                                           bufObj, cudaDeviceId);

  if (*image == NULL) {
    printf("Unable to create image\n");
    return NVMEDIA_STATUS_ERROR;
  }
  InitImage(*image, surfAllocAttrs[0].value, surfAllocAttrs[1].value);

  /*    printf("%s: NvMediaImageCreate:: Image size: %ux%u Image type: %d\n",
              __func__, surfAllocAttrs[0].value, surfAllocAttrs[1].value,
     surfType);*/

  return NVMEDIA_STATUS_OK;
}

/* Create NvMediaImage surface based on the input attributes.
 * Returns NVMEDIA_STATUS_OK on success
 */
static NvMediaStatus createSurfaceNonNvSCI(
    Blit2DTest *ctx, NvMediaSurfFormatAttr *surfFormatAttrs,
    NvMediaSurfAllocAttr *surfAllocAttrs, uint32_t numSurfAllocAttrs,
    NvMediaImage **image) {
  NvMediaSurfaceType surfType;

  /* create source image */
  surfType =
      NvMediaSurfaceFormatGetType(surfFormatAttrs, NVM_SURF_FMT_ATTR_MAX);

  *image = NvMediaImageCreateNew(ctx->device, surfType, surfAllocAttrs,
                                 numSurfAllocAttrs, 0);

  if (*image == NULL) {
    printf("Unable to create image\n");
    return NVMEDIA_STATUS_ERROR;
  }
  InitImage(*image, surfAllocAttrs[0].value, surfAllocAttrs[1].value);

  /*    printf("%s: NvMediaImageCreate:: Image size: %ux%u Image type: %d\n",
              __func__, surfAllocAttrs[0].value, surfAllocAttrs[1].value,
     surfType);*/

  return NVMEDIA_STATUS_OK;
}

static void destroySurface(NvMediaImage *image) { NvMediaImageDestroy(image); }

static NvMediaStatus blit2DImage(Blit2DTest *ctx, TestArgs *args,
                                 NvSciSyncObj &nvMediaSignalerSyncObj,
                                 NvSciSyncFence *preSyncFence,
                                 NvSciSyncFence *fence) {
  NvMediaStatus status;
  NvMediaImageSurfaceMap surfaceMap;

  status = ReadImage(args->inputFileName,              /* fileName */
                     0,                                /* frameNum */
                     args->srcSurfAllocAttrs[0].value, /* source image width */
                     args->srcSurfAllocAttrs[1].value, /* source image height */
                     ctx->srcImage,                    /* srcImage */
                     NVMEDIA_TRUE,                     /* uvOrderFlag */
                     1,                                /* bytesPerPixel */
                     MSB_ALIGNED);                     /* pixelAlignment */

  if (status != NVMEDIA_STATUS_OK) {
    printf("%s: ReadImage failed for input buffer: %d\n", __func__, status);
    return status;
  }

  if ((args->srcRect.x1 <= args->srcRect.x0) ||
      (args->srcRect.y1 <= args->srcRect.y0)) {
    ctx->srcRect = NULL;
  } else {
    ctx->srcRect = &(args->srcRect);
  }

  if ((args->dstRect.x1 <= args->dstRect.x0) ||
      (args->dstRect.y1 <= args->dstRect.y0)) {
    ctx->dstRect = NULL;
  } else {
    ctx->dstRect = &(args->dstRect);
  }

  static int64_t launch = 0;
  // Start inserting pre-fence from second launch inorder to for NvMedia2Blit to
  // wait
  // for cuda signal on fence.
  if (launch) {
    status = NvMedia2DInsertPreNvSciSyncFence(ctx->i2d, preSyncFence);
    if (status != NVMEDIA_STATUS_OK) {
      printf("%s: NvMedia2DSetNvSciSyncObjforEOF   failed: %d\n", __func__,
             status);
      return status;
    }
    NvSciSyncFenceClear(preSyncFence);
  }
  launch++;

  status = NvMedia2DSetNvSciSyncObjforEOF(ctx->i2d, nvMediaSignalerSyncObj);
  if (status != NVMEDIA_STATUS_OK) {
    printf("%s: NvMedia2DSetNvSciSyncObjforEOF   failed: %d\n", __func__,
           status);
    return status;
  }

  /* 2DBlit processing on input image */
  status = NvMedia2DBlitEx(ctx->i2d,          /* i2d */
                           ctx->dstImage,     /* dstSurface */
                           ctx->dstRect,      /* dstRect */
                           ctx->srcImage,     /* srcSurface */
                           ctx->srcRect,      /* srcRect */
                           &args->blitParams, /* params */
                           NULL);             /* paramsOut */

  if (status != NVMEDIA_STATUS_OK) {
    printf("%s: NvMedia2DBlitEx failed: %d\n", __func__, status);
    return status;
  }

  status =
      NvMedia2DGetEOFNvSciSyncFence(ctx->i2d, nvMediaSignalerSyncObj, fence);
  if (status != NVMEDIA_STATUS_OK) {
    printf("%s: NvMedia2DGetEOFNvSciSyncFence failed: %d\n", __func__, status);
    return status;
  }

  return NVMEDIA_STATUS_OK;
}

static NvMediaStatus blit2DImageNonNvSCI(Blit2DTest *ctx, TestArgs *args) {
  NvMediaStatus status;
  NvMediaImageSurfaceMap surfaceMap;

  status = ReadImage(args->inputFileName,              /* fileName */
                     0,                                /* frameNum */
                     args->srcSurfAllocAttrs[0].value, /* source image width */
                     args->srcSurfAllocAttrs[1].value, /* source image height */
                     ctx->srcImage,                    /* srcImage */
                     NVMEDIA_TRUE,                     /* uvOrderFlag */
                     1,                                /* bytesPerPixel */
                     MSB_ALIGNED);                     /* pixelAlignment */

  if (status != NVMEDIA_STATUS_OK) {
    printf("%s: ReadImage failed for input buffer: %d\n", __func__, status);
    return status;
  }

  if ((args->srcRect.x1 <= args->srcRect.x0) ||
      (args->srcRect.y1 <= args->srcRect.y0)) {
    ctx->srcRect = NULL;
  } else {
    ctx->srcRect = &(args->srcRect);
  }

  if ((args->dstRect.x1 <= args->dstRect.x0) ||
      (args->dstRect.y1 <= args->dstRect.y0)) {
    ctx->dstRect = NULL;
  } else {
    ctx->dstRect = &(args->dstRect);
  }

  /* 2DBlit processing on input image */
  status = NvMedia2DBlitEx(ctx->i2d,          /* i2d */
                           ctx->dstImage,     /* dstSurface */
                           ctx->dstRect,      /* dstRect */
                           ctx->srcImage,     /* srcSurface */
                           ctx->srcRect,      /* srcRect */
                           &args->blitParams, /* params */
                           NULL);             /* paramsOut */
  if (status != NVMEDIA_STATUS_OK) {
    printf("%s: NvMedia2DBlitEx failed: %d\n", __func__, status);
    return status;
  }

  /* Write output image into buffer */
  ctx->bytesPerPixel = 1;
  WriteImageToAllocatedBuffer(ctx, ctx->dstImage, NVMEDIA_TRUE, NVMEDIA_FALSE,
                              ctx->bytesPerPixel);

  return NVMEDIA_STATUS_OK;
}

static void cleanup(Blit2DTest *ctx, NvMediaStatus status = NVMEDIA_STATUS_OK) {
  if (ctx->srcImage != NULL) {
    NvMedia2DImageUnRegister(ctx->i2d, ctx->srcImage);
    destroySurface(ctx->srcImage);
  }
  if (ctx->dstImage != NULL) {
    NvMedia2DImageUnRegister(ctx->i2d, ctx->dstImage);
    destroySurface(ctx->dstImage);
  }
  if (status != NVMEDIA_STATUS_OK) {
    exit(EXIT_FAILURE);
  }
}

void cleanupNvMedia(Blit2DTest *ctx, NvSciSyncObj &syncObj,
                    NvSciSyncObj &preSyncObj) {
  NvMediaStatus status;
  cleanup(ctx);
  status = NvMedia2DUnregisterNvSciSyncObj(ctx->i2d, syncObj);
  if (status != NVMEDIA_STATUS_OK) {
    printf("%s: NvMediaImageSciBufInit failed\n", __func__);
    exit(EXIT_FAILURE);
  }
  status = NvMedia2DUnregisterNvSciSyncObj(ctx->i2d, preSyncObj);
  if (status != NVMEDIA_STATUS_OK) {
    printf("%s: NvMediaImageSciBufInit failed\n", __func__);
    exit(EXIT_FAILURE);
  }
  NvMediaImageNvSciBufDeinit();
}

void cleanupNvMedia(Blit2DTest *ctx) {
  cleanup(ctx);
  free(ctx->dstBuffPitches);
  free(ctx->dstBuffer);
  free(ctx->dstBuff);
}

void setupNvMedia(TestArgs *args, Blit2DTest *ctx, NvSciBufObj &srcNvSciBufobj,
                  NvSciBufObj &dstNvSciBufobj, NvSciSyncObj &syncObj,
                  NvSciSyncObj &preSyncObj, int cudaDeviceId) {
  NvMediaStatus status;
  status = NvMediaImageNvSciBufInit();
  if (status != NVMEDIA_STATUS_OK) {
    printf("%s: NvMediaImageSciBufInit failed\n", __func__);
    cleanup(ctx, status);
  }

  // Create source surface
  status = createSurface(ctx, args->srcSurfFormatAttrs, args->srcSurfAllocAttrs,
                         args->numSurfAllocAttrs, &ctx->srcImage,
                         srcNvSciBufobj, cudaDeviceId);
  if (status != NVMEDIA_STATUS_OK) {
    printf("%s: Unable to create buffer pools\n", __func__);
    cleanup(ctx, status);
  }

  // Create destination surface
  status = createSurface(ctx, args->dstSurfFormatAttrs, args->dstSurfAllocAttrs,
                         args->numSurfAllocAttrs, &ctx->dstImage,
                         dstNvSciBufobj, cudaDeviceId);
  if (status != NVMEDIA_STATUS_OK) {
    printf("%s: Unable to create buffer pools\n", __func__);
    cleanup(ctx, status);
  }

  // Register source  Surface
  status =
      NvMedia2DImageRegister(ctx->i2d, ctx->srcImage, NVMEDIA_ACCESS_MODE_READ);
  if (status != NVMEDIA_STATUS_OK) {
    printf("%s: Unable to register source surface\n", __func__);
    cleanup(ctx, status);
  }
  // Register destination Surface
  status = NvMedia2DImageRegister(ctx->i2d, ctx->dstImage,
                                  NVMEDIA_ACCESS_MODE_READ_WRITE);
  if (status != NVMEDIA_STATUS_OK) {
    printf("%s: Unable to register destination surface\n", __func__);
    cleanup(ctx, status);
  }

  status = NvMedia2DRegisterNvSciSyncObj(ctx->i2d, NVMEDIA_EOFSYNCOBJ, syncObj);
  if (status != NVMEDIA_STATUS_OK) {
    printf("%s: Unable to NvMedia2DRegisterNvSciSyncObj\n", __func__);
  }

  status =
      NvMedia2DRegisterNvSciSyncObj(ctx->i2d, NVMEDIA_PRESYNCOBJ, preSyncObj);
  if (status != NVMEDIA_STATUS_OK) {
    printf("%s: Unable to NvMedia2DRegisterNvSciSyncObj\n", __func__);
  }
}

// Create NvMedia src & dst image without NvSciBuf
void setupNvMedia(TestArgs *args, Blit2DTest *ctx) {
  NvMediaStatus status;

  // Create source surface
  status = createSurfaceNonNvSCI(ctx, args->srcSurfFormatAttrs,
                                 args->srcSurfAllocAttrs,
                                 args->numSurfAllocAttrs, &ctx->srcImage);
  if (status != NVMEDIA_STATUS_OK) {
    printf("%s: Unable to create buffer pools\n", __func__);
    cleanup(ctx, status);
  }

  // Create destination surface
  status = createSurfaceNonNvSCI(ctx, args->dstSurfFormatAttrs,
                                 args->dstSurfAllocAttrs,
                                 args->numSurfAllocAttrs, &ctx->dstImage);
  if (status != NVMEDIA_STATUS_OK) {
    printf("%s: Unable to create buffer pools\n", __func__);
    cleanup(ctx, status);
  }

  // Register source  Surface
  status =
      NvMedia2DImageRegister(ctx->i2d, ctx->srcImage, NVMEDIA_ACCESS_MODE_READ);
  if (status != NVMEDIA_STATUS_OK) {
    printf("%s: Unable to register source surface\n", __func__);
    cleanup(ctx, status);
  }

  // Register destination Surface
  status = NvMedia2DImageRegister(ctx->i2d, ctx->dstImage,
                                  NVMEDIA_ACCESS_MODE_READ_WRITE);
  if (status != NVMEDIA_STATUS_OK) {
    printf("%s: Unable to register destination surface\n", __func__);
    cleanup(ctx, status);
  }

  // Allocate buffer for writing image & set image parameters in Blit2DTest.
  ctx->bytesPerPixel = 1;
  AllocateBufferToWriteImage(ctx, ctx->dstImage, NVMEDIA_TRUE, /* uvOrderFlag */
                             NVMEDIA_FALSE);                   /* appendFlag */
}

void runNvMediaBlit2D(TestArgs *args, Blit2DTest *ctx) {
  // Blit2D function
  NvMediaStatus status = blit2DImageNonNvSCI(ctx, args);
  if (status != NVMEDIA_STATUS_OK) {
    printf("%s: Blit2D failed\n", __func__);
    cleanup(ctx, status);
  }
}

void runNvMediaBlit2D(TestArgs *args, Blit2DTest *ctx,
                      NvSciSyncObj &nvMediaSignalerSyncObj,
                      NvSciSyncFence *preSyncFence, NvSciSyncFence *fence) {
  // Blit2D function
  NvMediaStatus status =
      blit2DImage(ctx, args, nvMediaSignalerSyncObj, preSyncFence, fence);
  if (status != NVMEDIA_STATUS_OK) {
    printf("%s: Blit2D failed\n", __func__);
    cleanup(ctx, status);
  }
}