/*
 * Copyright (c) 2019, NVIDIA CORPORATION.  All rights reserved.
 *
 * NVIDIA CORPORATION and its licensors retain all intellectual property
 * and proprietary rights in and to this software, related documentation
 * and any modifications thereto.  Any use, reproduction, disclosure or
 * distribution of this software and related documentation without an express
 * license agreement from NVIDIA CORPORATION is strictly prohibited.
*/

/* standard headers */
#include <string.h>
#include <signal.h>

/* Nvidia headers */
#include "cmdline.h"
#include "log_utils.h"
#include "nvmedia_image.h"
#include "nvmedia_2d.h"
#include "nvmedia_surface.h"
#include "image_utils.h"
#include "image_scibuf_utils.h"
#include "nvmedia_image_nvscibuf.h"

#define CRC_GENERATION_MODE 1
#define CRC_CHECK_MODE 2

/* 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)
{
    NvMediaSurfaceType surfType;

    /* create source image */
    surfType = NvMediaSurfaceFormatGetType(surfFormatAttrs, NVM_SURF_FMT_ATTR_MAX);
    *image = NvMediaImageCreateUsingNvScibuf(ctx->device,             /* device */
                                             surfType,                /* surface type */
                                             surfAllocAttrs,
                                             numSurfAllocAttrs,
                                             0);
    if(*image == NULL) {
        LOG_ERR ("Unable to create image\n");
        return NVMEDIA_STATUS_ERROR;
    }
    InitImage (*image, surfAllocAttrs[0].value, surfAllocAttrs[1].value);

    LOG_DBG("%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)
{
    LOG_DBG("%s: Destroying surfaces\n", __func__);
    NvMediaImageDestroy(image);
}

static NvMediaStatus
blit2DImage(Blit2DTest *ctx, TestArgs* args)
{
    NvMediaStatus status;
    NvMediaImageSurfaceMap surfaceMap;
    uint64_t startTime,endTime;
    double processingTime;

    processingTime = 0;
    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) {
        LOG_ERR("%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);
    }

    GetTimeMicroSec(&startTime);

    /* 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) {
        LOG_ERR("%s: NvMedia2DBlitEx failed: %d\n", __func__, status);
        return status;
    }

    /* Write output image into file */
    status = NvMediaImageLock(ctx->dstImage,
                              NVMEDIA_IMAGE_ACCESS_WRITE,
                              &surfaceMap);
    if(status != NVMEDIA_STATUS_OK) {
        LOG_ERR("%s: NvMediaImageLock failed: %d\n", __func__, status);
        return status;
    }

    GetTimeMicroSec(&endTime);
    processingTime += (double)(endTime - startTime)/1000.0;
    NvMediaImageUnlock(ctx->dstImage);

    if (args->outputFileName) {
        LOG_DBG("WriteImage : Saving output image into file...\n");
        if(NVMEDIA_STATUS_OK !=
                WriteImage(args->outputFileName,
                           ctx->dstImage,
                           NVMEDIA_TRUE,        /* uvOrderFlag */
                           NVMEDIA_FALSE,       /* appendFlag */
                           1,                   /* bytesPerPixel */
                           NULL)) {
            LOG_ERR("%s: WriteImage  failed for output image : %d\n", __func__, status);
        }
    }
    LOG_INFO("Processing time per frame %.4f ms \n",processingTime);

    return NVMEDIA_STATUS_OK;
}

int main (int  argc, char *argv[])
{
    TestArgs args;
    Blit2DTest ctx;
    FILE *crcFile = NULL;
    uint32_t crcValue;
    NvMediaStatus status = NVMEDIA_STATUS_ERROR;

    /* Read configuration from command line and config file */
    memset(&args, 0, sizeof(TestArgs));
    memset(&ctx, 0, sizeof(Blit2DTest));

    /* ParseArgs parses the command line and the 2D configuration file and populates all initParams
     * and run time configuration in to appropriate structures within args
     */
    if (ParseArgs(argc, argv, &args)) {
        PrintUsage();
        return -1;
    }
    /* Check version */
    NvMediaVersion version;
    status = NvMedia2DGetVersion(&version);
    if (status == NVMEDIA_STATUS_OK) {
        LOG_INFO("Library version: %u.%u\n", version.major, version.minor);
        LOG_INFO("Header version:  %u.%u\n", NVMEDIA_2D_VERSION_MAJOR, NVMEDIA_2D_VERSION_MINOR);
        if ((version.major != NVMEDIA_2D_VERSION_MAJOR) ||
            (version.minor != NVMEDIA_2D_VERSION_MINOR)) {
            LOG_ERR("Library and Header mismatch!\n");
            goto end;
        }
    }

    // Create NvMedia device
    ctx.device = NvMediaDeviceCreate();
    if(!ctx.device) {
        LOG_ERR("%s: Failed to create NvMedia device\n", __func__);
        goto end;
    }

    // Create 2D blitter
    ctx.i2d = NvMedia2DCreate(ctx.device);
    if(!ctx.i2d) {
        LOG_ERR("%s: Failed to create NvMedia 2D i2d\n", __func__);
        goto end;
    }

    status = NvMediaImageNvSciBufInit();
    if(status != NVMEDIA_STATUS_OK) {
        LOG_ERR("%s: NvMediaImageSciBufInit failed\n",__func__);
        goto end;
    }

    // Create source surface
    status = createSurface(&ctx,
                           args.srcSurfFormatAttrs,
                           args.srcSurfAllocAttrs,
                           args.numSurfAllocAttrs,
                           &ctx.srcImage);
    if(status != NVMEDIA_STATUS_OK) {
        LOG_ERR("%s: Unable to create buffer pools\n", __func__);
        goto end;
    }

    // Create destination surface
    status = createSurface(&ctx,
                           args.dstSurfFormatAttrs,
                           args.dstSurfAllocAttrs,
                           args.numSurfAllocAttrs,
                           &ctx.dstImage);
    if(status != NVMEDIA_STATUS_OK) {
        LOG_ERR("%s: Unable to create buffer pools\n", __func__);
        goto end;
    }
    //Register source  Surface
    status = NvMedia2DImageRegister(ctx.i2d, ctx.srcImage, NVMEDIA_ACCESS_MODE_READ);
    if ( status != NVMEDIA_STATUS_OK) {
        LOG_ERR("%s: Unable to register source surface\n", __func__);
        goto end;
    }
    //Register destination Surface
    status = NvMedia2DImageRegister(ctx.i2d, ctx.dstImage, NVMEDIA_ACCESS_MODE_READ_WRITE);
    if ( status != NVMEDIA_STATUS_OK) {
        LOG_ERR("%s: Unable to register destination surface\n", __func__);
        goto end;
    }
    // Blit2D function
    status = blit2DImage(&ctx, &args);
    if(status != NVMEDIA_STATUS_OK) {
        LOG_ERR("%s: Blit2D failed\n", __func__);
        goto end;
    }

    crcValue = 0;
    if (args.crcMode == CRC_GENERATION_MODE) {
        crcFile = fopen(args.CRCFileName, "wt");
        if (!crcFile){
            LOG_ERR("Cannot open crc File\n");
            goto end;
        }
        status = GetImageCrc(ctx.dstImage,
                             args.dstSurfAllocAttrs[0].value,  /* destination image width */
                             args.dstSurfAllocAttrs[1].value,  /* destination image height */
                             &crcValue, 1);
        if (status != NVMEDIA_STATUS_OK) {
             LOG_ERR("GetImageCrc failed \n");
             goto end;
        }
        if (!fprintf(crcFile, "%08x\n", crcValue)) {
            LOG_ERR("main: Failed writing calculated CRC to file %s\n", args.CRCFileName);
        }
    } else if (args.crcMode == CRC_CHECK_MODE) {
        NvMediaBool isMatching;
        crcFile = fopen(args.CRCFileName, "rb");
        if(!crcFile){
            LOG_ERR("Cannot open crc File\n");
            goto end;
        }
        if (fscanf(crcFile, "%8x\n", &crcValue) == 1) {
            status = CheckImageCrc(ctx.dstImage,
                                   args.dstSurfAllocAttrs[0].value,  /* destination image width */
                                   args.dstSurfAllocAttrs[1].value,  /* destination image height */
                                   crcValue, &isMatching, 1);
            if (status != NVMEDIA_STATUS_OK) {
                LOG_ERR("CheckImageCrc failed \n");
                goto end;
            }
            if (!isMatching)
                LOG_MSG ("CRC Mismatch\n");
            else {
                LOG_MSG ("CRC Match with crcFileName: %s\n", args.CRCFileName);
            }
         } else {
            LOG_ERR("main: Failed reading file %s\n", args.CRCFileName);
         }
    }

end:
    if (crcFile) {
        fclose(crcFile);
    }

    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 (ctx.i2d != NULL) {
        NvMedia2DDestroy(ctx.i2d);
    }

    NvMediaImageNvSciBufDeinit();

    if (ctx.device != NULL) {
        NvMediaDeviceDestroy(ctx.device);
    }

    if (status == NVMEDIA_STATUS_OK) {
        return 0;
    }
    else {
        return 1;
    }
}
