/* Copyright (c) 2018-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.
 */

#include <stdio.h>
#include <string.h>
#include <nvmedia_core.h>
#include <log_utils.h>
#include <nvmedia_2d.h>
#include <cmdline.h>
#include "image_utils.h"

int main(int argc, char *argv[])
{
    NvMediaStatus status = NVMEDIA_STATUS_OK;

    TestArgs testArgs;
    int parseArgsStatus;

    NvMediaImage *inputImage1=NULL, *inputImage2=NULL;
    NvMediaImage *outImage=NULL;
    NvMediaDevice *device;
    NvMedia2D *i2d;
    NvMediaSurfAllocAttr surfAllocAttrs[7];
    NvMediaSurfaceType surfaceType = NvMediaSurfaceType_Unsupported;
    NvMediaSurfaceType outSurfaceType = NvMediaSurfaceType_Unsupported;
    uint32_t numAllocParam=0;
    uint32_t outWidth=0, outHeight=0;

    NVM_SURF_FMT_DEFINE_ATTR(surfFormatAttrs);
    NVM_SURF_FMT_DEFINE_ATTR(outSurfFormatAttrs);

    memset(&testArgs, 0, sizeof(TestArgs));

    parseArgsStatus = ParseArgs(argc, argv, &testArgs);
    if (parseArgsStatus) {
        if (parseArgsStatus != 1) {
            LOG_ERR("Failed to parse args\n");
            PrintUsage();
        }
        return -1;
    }

    if (testArgs.inputSurfType == NVM_SURF_ATTR_SURF_TYPE_RGBA ) {
        NVM_SURF_FMT_SET_ATTR_RGBA(surfFormatAttrs, RGBA, UINT, 8, PL);
    } else if (testArgs.inputSurfType == NVM_SURF_ATTR_SURF_TYPE_RAW) {
        NVM_SURF_FMT_SET_ATTR_RAW(surfFormatAttrs, RGGB, INT, 8, PL);
    }
    surfaceType = NvMediaSurfaceFormatGetType(surfFormatAttrs, NVM_SURF_FMT_ATTR_MAX);

    surfAllocAttrs[0].type = NVM_SURF_ATTR_WIDTH;
    surfAllocAttrs[0].value = testArgs.inputWidth;
    surfAllocAttrs[1].type = NVM_SURF_ATTR_HEIGHT;
    surfAllocAttrs[1].value = testArgs.inputHeight;
    numAllocParam = 2;

    device = NvMediaDeviceCreate();
    if (!device) {
        LOG_ERR("Failed to create nvmedia device\n");
        return -1;
    }

    i2d = NvMedia2DCreate(device);
    if (!i2d) {
        LOG_ERR("Failed to create 2D device\n");
        return -1;
    }

    inputImage1 = NvMediaImageCreateNew(device,
                                        surfaceType,
                                        surfAllocAttrs,
                                        numAllocParam,
                                        0);
    if (!inputImage1) {
        LOG_ERR("Failed to create inVideoSurface\n");
        return -1;
    }

    inputImage2 = NvMediaImageCreateNew(device,
                                        surfaceType,
                                        surfAllocAttrs,
                                        numAllocParam,
                                        0);
    if (!inputImage2) {
        LOG_ERR("Failed to create inVideoSurface\n");
        return -1;
    }

    NVM_SURF_FMT_SET_ATTR_RGBA(outSurfFormatAttrs, RGBA, UINT, 8, PL);
    outSurfaceType = NvMediaSurfaceFormatGetType(outSurfFormatAttrs, 7);

    if (testArgs.inputSurfType == NVM_SURF_ATTR_SURF_TYPE_RGBA) {
       outWidth = testArgs.inputWidth;
       outHeight = testArgs.inputHeight*2;
    } else if (testArgs.inputSurfType == NVM_SURF_ATTR_SURF_TYPE_RAW) {
       outWidth = testArgs.inputWidth/4;
       outHeight = testArgs.inputHeight*2;
    }
    surfAllocAttrs[0].type = NVM_SURF_ATTR_WIDTH;
    surfAllocAttrs[0].value = outWidth;
    surfAllocAttrs[1].type = NVM_SURF_ATTR_HEIGHT;
    surfAllocAttrs[1].value = outHeight;
    numAllocParam = 2;

    outImage = NvMediaImageCreateNew(device,
                                     outSurfaceType,
                                     surfAllocAttrs,
                                     numAllocParam,
                                     0);
    if (!outImage) {
        LOG_ERR("Failed to create outVideoSurface2\n");
        return -1;
    }

    status = ReadImage(testArgs.inputFileName1,
                       0,
                       testArgs.inputWidth,
                       testArgs.inputHeight,
                       inputImage1,
                       0,
                       0,
                       MSB_ALIGNED);
    if (status != NVMEDIA_STATUS_OK) {
        LOG_ERR("Failed to read 1st input frame (%d)\n", status);
        return status;
    }
    status = ReadImage(testArgs.inputFileName2,
                       0,
                       testArgs.inputWidth,
                       testArgs.inputHeight,
                       inputImage2,
                       0,
                       0,
                       MSB_ALIGNED);
    if (status != NVMEDIA_STATUS_OK) {
        LOG_ERR("Failed to read 2nd input frame (%d)\n", status);
        return status;
    }

#ifdef DEBUG_DUMP
    status = WriteImage("testInput1.rgba",
                        inputImage1,
                        NVMEDIA_FALSE,
                        NVMEDIA_FALSE,
                        4,
                        NULL);
    if (status != NVMEDIA_STATUS_OK) {
        LOG_ERR("Failed to write weave output\n");
        return status;
    }
    status = WriteImage("testInput2.rgba",
                        inputImage2,
                        NVMEDIA_FALSE,
                        NVMEDIA_FALSE,
                        4,
                        NULL);
    if (status != NVMEDIA_STATUS_OK) {
        LOG_ERR("Failed to write weave output\n");
        return status;
    }
#endif

    status = NvMedia2DWeaveNew(i2d,
                               inputImage1,
                               inputImage2,
                               outImage);
    if (status != NVMEDIA_STATUS_OK) {
            LOG_ERR("Failed to weave surfaces\n");
            return status;
    }

    if (testArgs.inputSurfType == NVM_SURF_ATTR_SURF_TYPE_RAW) {
        status = WriteImage(testArgs.outputFileName,
                            outImage,
                            NVMEDIA_FALSE,
                            NVMEDIA_FALSE,
                            1,
                            NULL);
    } else {
        status = WriteImage(testArgs.outputFileName,
                            outImage,
                            NVMEDIA_FALSE,
                            NVMEDIA_FALSE,
                            4,
                            NULL);
    }
    if (status != NVMEDIA_STATUS_OK) {
        LOG_ERR("Failed to write weave output\n");
        return status;
    }

    NvMediaImageDestroy(inputImage1);
    NvMediaImageDestroy(inputImage2);
    NvMediaImageDestroy(outImage);
    NvMedia2DDestroy(i2d);
    NvMediaDeviceDestroy(device);

    return 0;
}
