/*
 * Copyright (c) 2014-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 <stdlib.h>
#include <stdbool.h>
#include <string.h>
#include <sys/time.h>
#include <sys/stat.h>
#include <unistd.h>
#include <signal.h>

#include "cmdline.h"
#include "config_parser.h"
#include "log_utils.h"
#include "misc_utils.h"
#include "nvmedia_image.h"
#include "nvmedia_ijpe.h"
#include "nvmedia_surface.h"
#include "image_utils.h"

static NvMediaBool encodeStop = 0;

#define MAX_BITSTREAM_SIZE (10 * 1024 * 1024)
#define MAX_HUFFMAN_CODES        256
#define MAX_HUFFMAN_CODE_LENGTH  16
#define QUANTIZATION_TABLE_SIZE  64

/* Signal Handler for SIGINT */
static void sigintHandler(int sig_num)
{
    LOG_MSG("\n Exiting encode process \n");
    encodeStop = 1;
}

static NvMediaStatus
CheckVersion(void)
{
    NvMediaVersion version;
    NvMediaStatus status = NVMEDIA_STATUS_OK;

    memset(&version, 0, sizeof(NvMediaVersion));
    status = NvMediaCoreGetVersion(&version);
    if (status != NVMEDIA_STATUS_OK)
        return status;

    if((version.major != NVMEDIA_CORE_VERSION_MAJOR) ||
       (version.minor != NVMEDIA_CORE_VERSION_MINOR)) {
        LOG_ERR("%s: Incompatible core version found \n", __func__);
        LOG_ERR("%s: Client version: %d.%d\n", __func__,
            NVMEDIA_CORE_VERSION_MAJOR, NVMEDIA_CORE_VERSION_MINOR);
        LOG_ERR("%s: Core version: %d.%d\n", __func__,
            version.major, version.minor);
        return NVMEDIA_STATUS_INCOMPATIBLE_VERSION;
    }

    memset(&version, 0, sizeof(NvMediaVersion));
    status = NvMediaImageGetVersion(&version);
    if (status != NVMEDIA_STATUS_OK)
        return status;

    if((version.major != NVMEDIA_IMAGE_VERSION_MAJOR) ||
       (version.minor != NVMEDIA_IMAGE_VERSION_MINOR)) {
        LOG_ERR("%s: Incompatible image version found \n", __func__);
        LOG_ERR("%s: Client version: %d.%d\n", __func__,
            NVMEDIA_IMAGE_VERSION_MAJOR, NVMEDIA_IMAGE_VERSION_MINOR);
        LOG_ERR("%s: Core version: %d.%d\n", __func__,
            version.major, version.minor);
        return NVMEDIA_STATUS_INCOMPATIBLE_VERSION;
    }

    memset(&version, 0, sizeof(NvMediaVersion));
    status = NvMediaIJPEGetVersion(&version);
    if (status != NVMEDIA_STATUS_OK)
        return status;

    if((version.major != NVMEDIA_IJPE_VERSION_MAJOR) ||
       (version.minor != NVMEDIA_IJPE_VERSION_MINOR)) {
        LOG_ERR("%s: Incompatible JPEG encode version found \n", __func__);
        LOG_ERR("%s: Client version: %d.%d\n", __func__,
            NVMEDIA_IJPE_VERSION_MAJOR, NVMEDIA_IJPE_VERSION_MINOR);
        LOG_ERR("%s: Core version: %d.%d\n", __func__,
            version.major, version.minor);
        return NVMEDIA_STATUS_INCOMPATIBLE_VERSION;
    }

    return status;
}

/*******************************************************************************
*   Function Name: readInputTable
*   Purpose:
*       Search for the input string "pString" in the input file
*       and reads "numElements" number of values from that point
*       and store it in the location - pTable
*
*   Inputs:
*        fp          - input file
*        pString     - String to be searched
*        pTable      - Location to save read values
*        numElements - number of elements to be read
*
*   Return Data:
*       Status       - 1 on success and 0 on failure.
*
*******************************************************************************/
static uint32_t
readInputTable(FILE *fp, const char *pString, unsigned char *pTable, uint32_t numElements)
{
    char temp[256];
    uint32_t j;
    bool found = false;
    while ((fscanf(fp, "%s", temp)) != EOF) {
        if (strcmp(temp, pString) == 0){
            found = true;
            break;
        }
    }

    if (found == false) {
        LOG_ERR("Huffman/Quant table is not found\n");
        return 0;
    }

    if ((fscanf(fp, "%s %s", temp + 10, temp + 20)) == EOF) {
        LOG_ERR("Error in huffman Table entry format\n");
        return 0;
    }

    /* Reading  values*/
    for (j = 0; j < numElements; j++) {
        if ((fscanf(fp, "%hhu,", &pTable[j])) == EOF) {
            LOG_ERR("Error in huffman Table entry format\n");
            return 0;
        }
    }
    return 1;
}

/*******************************************************************************
*   Function Name: sumOfTable
*   Purpose:
*       Add first n elements of the input array and returns the sum
*
*   Inputs:
*        array       - Input array
*        numElements - number of elements to be added
*
*   Return Data:
*        sum       - sum of the elements
*
*******************************************************************************/
static uint32_t
sumOfTable(unsigned char* array, uint32_t numElements)
{
    uint32_t sum = 0; // initialize sum
    uint32_t i;

    // Iterate through all elements and add them to sum
    for (i = 0; i < numElements; i++) {
        sum += array[i];
    }

    return sum;
}

static uint32_t
readHuffmanTable(FILE *huffFile,
                 NvMediaJPHuffmanTableSpecfication *colorComponent,
                 const char *pStringBits,
                 const char *pStringVal)
{
    uint32_t numValues;

    /* reads number of codes with length 1-16 */
    if (!readInputTable(huffFile, pStringBits, colorComponent->length, MAX_HUFFMAN_CODE_LENGTH)) {
        return 0;
    }

    /* calculate total number of codes */
    numValues = sumOfTable(colorComponent->length, MAX_HUFFMAN_CODE_LENGTH);
    /* total number of codes should not exceed 256 */
    if (numValues > MAX_HUFFMAN_CODES) {
        LOG_ERR("Error in Huffman table\n");
        return 0;
    }

    colorComponent->values = calloc(MAX_HUFFMAN_CODES, sizeof(unsigned char));
    if (!colorComponent->values) {
        LOG_ERR("Failed to allocate memory for huffman code values\n");
        return 0;
    }
    /* get values for each code */
    return (readInputTable(huffFile, pStringVal, colorComponent->values, numValues));
}

int main(int argc, char *argv[])
{
    TestArgs args;
    FILE *crcFile = NULL, *outputFile = NULL, *streamFile = NULL;
    FILE *huffFile = NULL, *quantFile = NULL;
    char outFileName[FILE_NAME_SIZE];
    NvMediaStatus status = NVMEDIA_STATUS_OK;
    NvMediaDevice *device;
    NvMediaImage *imageFrame = NULL;
    NvMediaIJPE *testEncoder = NULL;
    NvMediaBool nextFrameFlag = NVMEDIA_TRUE, encodeDoneFlag;
    long long totalBytes = 0;
    long fileLength;
    uint8_t *buffer = NULL;
    uint32_t framesNum = 0, frameCounter = 0, bytes, bytesAvailable = 0, calcCrc = 0;
    uint32_t imageSize = 0;
    uint64_t startTime, endTime1, endTime2;
    double encodeTime = 0;
    double getbitsTime = 0;
    NvMediaSurfAllocAttr surfAllocAttrs[8];
    uint32_t numSurfAllocAttrs = 0;
    NvMediaJPEncAttributes attr;
    uint32_t i;
    NvMediaBool testPass = NVMEDIA_FALSE;

    signal(SIGINT, sigintHandler);
    signal(SIGTERM, sigintHandler);

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

    LOG_DBG("main: Parsing jpeg encode command\n");
    if(!ParseArgs(argc, argv, &args)) {
        LOG_ERR("main: Parsing arguments failed\n");
        return -1;
    }

    if(CheckVersion() != NVMEDIA_STATUS_OK) {
        return -1;
    }

    if(args.crcoption.crcGenMode && args.crcoption.crcCheckMode) {
        LOG_ERR("main: crcGenMode and crcCheckMode cannot be enabled at the same time\n");
        return -1;
    }

    imageSize = (args.inputWidth * args.inputHeight * 3) / 2;

    LOG_DBG("main: NvMediaDeviceCreate\n");
    device = NvMediaDeviceCreate();
    if(!device) {
        LOG_ERR("main: NvMediaDeviceCreate failed\n");
        return -1;
    }

    LOG_DBG("main: Encode start from frame %d, imageSize=%d\n", frameCounter, imageSize);

    streamFile = fopen(args.infile, "rb");
    if(!streamFile) {
        LOG_ERR("main: Error opening '%s' for reading\n", args.infile);
        goto fail;
    }
    fseek(streamFile, 0, SEEK_END);
    fileLength = ftell(streamFile);
    fclose(streamFile);
    if(!fileLength) {
       LOG_ERR("main: Zero file length for file %s, len=%d\n", args.infile, (int)fileLength);
       goto fail;
    }
    framesNum = fileLength / imageSize;
    surfAllocAttrs[0].type = NVM_SURF_ATTR_WIDTH;
    surfAllocAttrs[0].value = args.inputWidth;
    surfAllocAttrs[1].type = NVM_SURF_ATTR_HEIGHT;
    surfAllocAttrs[1].value = args.inputHeight;
    surfAllocAttrs[2].type = NVM_SURF_ATTR_CPU_ACCESS;
    surfAllocAttrs[2].value = NVM_SURF_ATTR_CPU_ACCESS_UNCACHED;
    numSurfAllocAttrs = 3;

    imageFrame =  NvMediaImageCreateNew(device,                            // device
                                       args.inputSurfType,                // surface type
                                       surfAllocAttrs,
                                       numSurfAllocAttrs,
                                       0);         // config
    if(!imageFrame) {
        LOG_ERR("main: NvMediaImageCreate failed\n");
        goto fail;
    }

    testEncoder = NvMediaIJPECreate(device,
                                   args.inputSurfType,               // inputFormat
                                   args.maxOutputBuffering,          // maxOutputBuffering
                                   MAX_BITSTREAM_SIZE);              // maxBitstreamBytes
    if(!testEncoder) {
        LOG_ERR("main: NvMediaIJPECreate failed\n");
        goto fail;
    }

    LOG_DBG("main: NvMediaIJPECreate, %p\n", testEncoder);

    if(args.crcoption.crcGenMode){
        crcFile = fopen(args.crcoption.crcFilename, "wt");
        if(!crcFile){
            LOG_ERR("main: Cannot open crc gen file for writing\n");
            goto fail;
        }
    } else if(args.crcoption.crcCheckMode){
        crcFile = fopen(args.crcoption.crcFilename, "rb");
        if(!crcFile){
            LOG_ERR("main: Cannot open crc gen file for reading\n");
            goto fail;
        }
    }

    if (args.huffTable) {
        attr.lumaDC = malloc(sizeof(NvMediaJPHuffmanTableSpecfication));
        attr.chromaDC = malloc(sizeof(NvMediaJPHuffmanTableSpecfication));
        attr.lumaAC = malloc(sizeof(NvMediaJPHuffmanTableSpecfication));
        attr.chromaAC = malloc(sizeof(NvMediaJPHuffmanTableSpecfication));
        if ((!attr.lumaDC) || (!attr.chromaDC) || (!attr.lumaAC) || (!attr.chromaAC)) {
            LOG_ERR("Failed to allocate memory\n");
            goto fail;
        }

        huffFile = fopen(args.huffFileName, "r");
        if(!huffFile){
            LOG_ERR("main: Cannot open huffman-table file for reading\n");
            goto fail;
        }


        if (!(readHuffmanTable(huffFile,
                              attr.lumaDC,
                              "dc_luminance_bits[16]",
                              "dc_luminance_value"))) {
            LOG_ERR("Error in Huffman table\n");
            goto fail;
        }

        if (!(readHuffmanTable(huffFile,
                              attr.chromaDC,
                              "dc_chrominance_bits[16]",
                              "dc_chrominance_value"))) {
            LOG_ERR("Error in Huffman table\n");
            goto fail;
        }

        if (!(readHuffmanTable(huffFile,
                              attr.lumaAC,
                              "ac_luminance_bits[16]",
                              "ac_luminance_value"))) {
            LOG_ERR("Error in Huffman table\n");
            goto fail;
        }

        if (!(readHuffmanTable(huffFile,
                              attr.chromaAC,
                              "ac_chrominance_bits[16]",
                              "ac_chrominance_value"))) {
            LOG_ERR("Error in Huffman table\n");
            goto fail;
        }

        status = NvMediaIJPESetAttributes(testEncoder,
                                          NVMEDIA_IMAGE_JPEG_ATTRIBUTE_HUFFTABLE,
                                          &attr);
        if (status != NVMEDIA_STATUS_OK) {
           LOG_ERR("main: NvMediaIJPESetAttributes failed\n");
           goto fail;
        }
    }

    if (args.quantTable) {
        quantFile = fopen(args.quantFileName, "r");
        if (!quantFile){
            LOG_ERR("main: Cannot open quantization-table file for reading\n");
            goto fail;
        }

        if (!readInputTable(quantFile,
                            "luma_quant_tbl[64]",
                            attr.lumaQuant,
                            QUANTIZATION_TABLE_SIZE)) {
            LOG_ERR("Error in quantization table\n");
            goto fail;
        }

        if (!readInputTable(quantFile,
                            "chroma_quant_tbl[64]",
                            attr.chromaQuant,
                            QUANTIZATION_TABLE_SIZE)) {
            LOG_ERR("Error in quantization table\n");
            goto fail;
        }

        for (i = 0; i < QUANTIZATION_TABLE_SIZE; i++) {
            if ((attr.lumaQuant[i] == 0) || (attr.chromaQuant[i] == 0)) // invalid quant value
            {
                LOG_ERR("main: Invalid quant value\n");
                goto fail;
            }
        }

        status = NvMediaIJPESetAttributes(testEncoder,
                                          NVMEDIA_IMAGE_JPEG_ATTRIBUTE_QUANTTABLE,
                                          &attr);

        if (status != NVMEDIA_STATUS_OK) {
           LOG_ERR("main: NvMediaIJPESetAttributes failed\n");
           goto fail;
        }
    }


    while(nextFrameFlag && (encodeStop == 0u)) {
        // Read Frame
        LOG_DBG("main: Reading YUV frame %d from file %s to image surface location: %p. (W:%d, H:%d)\n",
                 frameCounter, args.infile, imageFrame, args.inputWidth, args.inputHeight);
        status = ReadImage(args.infile,
                           frameCounter,
                           args.inputWidth,
                           args.inputHeight,
                           imageFrame,
                           NVMEDIA_TRUE,            //inputUVOrderFlag
                           1,                       //rawBytesPerPixel
                           MSB_ALIGNED);            //pixelAlignment
        if(status != NVMEDIA_STATUS_OK) {
           LOG_ERR("readYUVFile failed\n");
           goto fail;
        }
        LOG_DBG("main: ReadYUVFrame %d/%d done\n", frameCounter, framesNum-1);

        GetTimeMicroSec(&startTime);
        LOG_DBG("main: Encoding frame #%d\n", frameCounter);

        if (args.quantTable) {
            status = NvMediaIJPEFeedFrameQuant(testEncoder,
                                               imageFrame,
                                               attr.lumaQuant,
                                               attr.chromaQuant);
        }
        else {
            status = NvMediaIJPEFeedFrame(testEncoder,
                                          imageFrame,
                                          args.quality);
        }

        if(status != NVMEDIA_STATUS_OK) {
            LOG_ERR("main: NvMediaIJPEFeedFrameQuality failed: %x\n", status);
            goto fail;
        }

        encodeDoneFlag = NVMEDIA_FALSE;
        while(!encodeDoneFlag) {
            bytesAvailable = 0;
            bytes = 0;
            status = NvMediaIJPEBitsAvailable(testEncoder,
                                             &bytesAvailable,
                                             NVMEDIA_ENCODE_BLOCKING_TYPE_IF_PENDING,
                                             NVMEDIA_VIDEO_ENCODER_TIMEOUT_INFINITE);
            switch(status) {
                case NVMEDIA_STATUS_OK:
                    // Encode Time
                    GetTimeMicroSec(&endTime1);
                    encodeTime += (double)(endTime1 - startTime) / 1000.0;

                    buffer = malloc(bytesAvailable);
                    if(!buffer) {
                        LOG_ERR("main: Error allocating %d bytes\n", bytesAvailable);
                        goto fail;
                    }
                    status = NvMediaIJPEGetBits(testEncoder, &bytes, buffer, 0);
                    if(status != NVMEDIA_STATUS_OK && status != NVMEDIA_STATUS_NONE_PENDING) {
                        LOG_ERR("main: Error getting encoded bits\n");
                        goto fail;
                    }

                    if(bytes != bytesAvailable) {
                        LOG_ERR("main: byte counts do not match %d vs. %d\n", bytesAvailable, bytes);
                        goto fail;
                    }

                    GetTimeMicroSec(&endTime2);
                    getbitsTime += (double)(endTime2 - endTime1) / 1000.0;

                    LOG_DBG("main: Opening output file\n");
                    sprintf(outFileName, args.outfile, frameCounter);
                    outputFile = fopen(outFileName, "w+");
                    if(!outputFile) {
                        LOG_ERR("main: Failed opening '%s' file for writing\n", args.outfile);
                        goto fail;
                    }

                    if(fwrite(buffer, bytesAvailable, 1, outputFile) != 1) {
                       LOG_ERR("main: Error writing %d bytes\n", bytesAvailable);
                       fclose(outputFile);
                       goto fail;
                    }
                    fclose(outputFile);

                    if(args.crcoption.crcGenMode){
                        //calculate CRC from buffer 'buffer'
                        calcCrc = 0;
                        calcCrc = CalculateBufferCRC(bytesAvailable, calcCrc, buffer);
                        if(!fprintf(crcFile, "%08x\n",calcCrc)) {
                            LOG_ERR("main: Failed writing calculated CRC to file %s\n", crcFile);
                            goto fail;
                        }
                    } else if(args.crcoption.crcCheckMode){
                        //calculate CRC from buffer 'buffer'
                        uint32_t refCrc;
                        calcCrc = 0;
                        calcCrc = CalculateBufferCRC(bytesAvailable, calcCrc, buffer);
                        if(fscanf(crcFile, "%8x\n", &refCrc) == 1) {
                            if(refCrc != calcCrc){
                                LOG_ERR("main: Frame %d crc 0x%x does not match with ref crc 0x%x\n",
                                        frameCounter, calcCrc, refCrc);
                                goto fail;
                            }
                        } else {
                            LOG_ERR("main: Failed checking CRC. Failed reading file %s\n", crcFile);
                            goto fail;
                        }
                    }

                    free(buffer);
                    buffer = NULL;

                    //Tracking the bitrate
                    totalBytes += bytesAvailable;

                    encodeDoneFlag = 1;
                    break;
                case NVMEDIA_STATUS_PENDING:
                    LOG_DBG("main: Status - pending\n");
                    break;
                case NVMEDIA_STATUS_NONE_PENDING:
                    LOG_ERR("main: No encoded data is pending\n");
                    goto fail;
                default:
                    LOG_ERR("main: Error occured\n");
                    goto fail;
            }
        }

        // Next frame
        frameCounter++;

        if(frameCounter == framesNum) {
            nextFrameFlag = NVMEDIA_FALSE;
        }
    }

    //get encoding time info
    LOG_MSG("\nTotal Encoding time for %d frames: %.3f ms\n", frameCounter, encodeTime + getbitsTime);
    LOG_MSG("Encoding time per frame %.4f ms \n", encodeTime / frameCounter);
    LOG_MSG("Get bits time per frame %.4f ms \n", getbitsTime / frameCounter);
    //Get the bitrate info
    LOG_MSG("\nTotal encoded frames = %d, avg. bitrate=%d\n",
            frameCounter, (int)(totalBytes*8*30/frameCounter));
    if (args.crcoption.crcGenMode){
        LOG_MSG("\n***crc gold file %s has been generated***\n", args.crcoption.crcFilename);
    } else if (args.crcoption.crcCheckMode){
        LOG_MSG("\n***crc checking with file %s is successful\n", args.crcoption.crcFilename);
    }
    LOG_MSG("\n***ENCODING PROCESS ENDED SUCCESSFULY***\n");
    testPass = NVMEDIA_TRUE;

fail:
    if(imageFrame) {
       NvMediaImageDestroy(imageFrame);
    }

    if(testEncoder) {
        NvMediaIJPEDestroy(testEncoder);
    }

    if(device) {
        LOG_DBG("main: Destroying device\n");
        NvMediaDeviceDestroy(device);
    }

    if(crcFile) {
        fclose(crcFile);
    }

    if (buffer)
        free(buffer);

    if (args.huffTable) {
        if (attr.lumaDC) {
            if (attr.lumaDC->values)
                free(attr.lumaDC->values);
            free(attr.lumaDC);
        }

        if (attr.chromaDC) {
            if (attr.chromaDC->values)
                free(attr.chromaDC->values);
            free(attr.chromaDC);
        }

        if (attr.lumaAC) {
            if (attr.lumaAC->values)
                free(attr.lumaAC->values);
            free(attr.lumaAC);
        }
        if (attr.chromaAC) {
            if (attr.chromaAC->values)
                free(attr.chromaAC->values);
            free(attr.chromaAC);
        }
    }

    if (testPass) {
        return 0;
    } else {
        return -1;
    }
}
