/*
 * Copyright (c) 2016-2018, 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 <string.h>
#include "img_producer.h"

extern volatile bool signal_stop;

static uint32_t GetFramesCount(char* fileName, int imageSize)
{
    uint32_t count = 0;
    FILE *inputFile = NULL;

    if(imageSize == 0) {
       LOG_ERR("%s: Bad image parameters", __func__);
       return 0;
    }

    if(!(inputFile = fopen(fileName, "rb"))) {
       LOG_ERR("%s: Failed opening file %s", __func__,
               fileName);
       goto done;
    }

    fseek(inputFile, 0, SEEK_END);
    count = ftell(inputFile) / imageSize;

    fclose(inputFile);

done:
    return count;
}

static void
*ProcessImageThread(void * args)
{
    ImageProducerCtx *imgProducer = (ImageProducerCtx *)args;
    NvMediaImage *inputImage = NULL;
    NvMediaImage *releaseImage = NULL;
    NvMediaStatus status;
    uint32_t readFrame = 0;
    uint32_t BytesPerPixel = 1;
    NvMediaTime timeStamp ={0, 0};
    NvMediaImage *filereadin = NULL;
    NVM_SURF_FMT_DEFINE_ATTR(attr);

    status = NvMediaSurfaceFormatGetAttrs(imgProducer->surfaceType,
                                          attr,
                                          NVM_SURF_FMT_ATTR_MAX);
    if (status != NVMEDIA_STATUS_OK) {
        LOG_ERR("%s:NvMediaSurfaceFormatGetAttrs failed\n", __func__);
        goto done;
    }

    if(attr[NVM_SURF_ATTR_SURF_TYPE].value == NVM_SURF_ATTR_SURF_TYPE_RAW){
        BytesPerPixel = 2;
    }
    while(!signal_stop) {

        /* acquire Image from inputQueue */
        if(IsFailed(NvQueueGet(imgProducer->inputQueue,
                               (void *)&inputImage,
                               100))) {
            LOG_INFO("%s: Image Producer input queue empty\n", __func__);
            goto getImage;
        }

        //Read Image into input queue
        if((imgProducer->frameCount && (readFrame+1) > imgProducer->frameCount)) {
            LOG_DBG("Image Producer read image done\n");
            goto done;
        }

        LOG_DBG("Image Producer: Reading image %u/%u\n", readFrame+1, imgProducer->frameCount);
        filereadin = imgProducer->useblitpath ? imgProducer->fileimage :
                        inputImage;
        if(IsFailed(ReadImage(imgProducer->inpFileName,
                              readFrame,
                              imgProducer->width,
                              imgProducer->height,
                              filereadin,
                              NVMEDIA_TRUE,                 //inputUVOrderFlag,
                              BytesPerPixel,                //rawInputBytesPerPixel
                              MSB_ALIGNED))) {              // pixelAlignment
            LOG_ERR("Image Producer: Failed reading frame %u", readFrame);
            NvQueuePut((NvQueue *)inputImage->tag,
                       (void *)&inputImage,
                       ENQUEUE_TIMEOUT);
            goto done;
        }

        readFrame++;

        if(imgProducer->useblitpath) {
            if(IsFailed(NvMedia2DBlitEx(imgProducer->blitter,
                                        inputImage,
                                        NULL,
                                        filereadin,
                                        NULL,
                                        NULL,
                                        NULL))) {
                LOG_ERR("%s: image NvMedia2DBlitEx failed \n",__func__);
                goto done;
            }
        }

        // Post image to egl-stream
        status = NvMediaEglStreamProducerPostImage(imgProducer->producer,
                                                   inputImage,
                                                   &timeStamp);

        if(status != NVMEDIA_STATUS_OK) {
            LOG_ERR("Image Producer: NvMediaEglStreamProducerPostImage failed, status=%d\n", status);
            goto done;
        }

getImage:
        // Get back from the egl-stream
        status = NvMediaEglStreamProducerGetImage(imgProducer->producer,
                                                  &releaseImage,
                                                  100);
        if (status == NVMEDIA_STATUS_OK) {
                    NvQueuePut((NvQueue *)releaseImage->tag,
                               (void *)&releaseImage,
                               ENQUEUE_TIMEOUT);
         } else {
            LOG_DBG ("Image Producer: NvMediaEglStreamProducerGetImage waiting\n");
            continue;
         }

    }

done:
    usleep(2000000);
    *imgProducer->producerDone = true;

    return NULL;
}

int ImageProducerInit (volatile bool *producerDone,
                      ImageProducerCtx *imgProducer,
                      EglStreamClient *streamClient,
                      TestArgs *args)
{
    NvMediaImage *image = NULL;
    uint32_t i = 0;
    NvMediaSurfAllocAttr surfAllocAttrs[8];
    uint32_t numSurfAllocAttrs = 0;
    uint32_t imageSize = 0;
    char *env;
    memset(imgProducer, 0, sizeof(ImageProducerCtx));
    imgProducer->inpFileName    = args->inpFileName;
    imgProducer->width          = args->width;
    imgProducer->height         = args->height;
    imgProducer->frameCount     = args->frameCount;
    imgProducer->display        = streamClient->display;
    imgProducer->eglStream      = streamClient->eglStream;
    imgProducer->producerDone   = producerDone;
    imgProducer->useblitpath = args->useblitpath;

    /* Define an array of NvMediaSurfFormatAttr */
    NVM_SURF_FMT_DEFINE_ATTR(surfFormatAttrs);

    /* set producer surface Format*/
    if(args->imagetype == IMAGE_TYPE_RGBA) {
        if(!args->pitchLinearOutput){
            NVM_SURF_FMT_SET_ATTR_RGBA(surfFormatAttrs, RGBA, UINT, 8, BL);
        } else {
            NVM_SURF_FMT_SET_ATTR_RGBA(surfFormatAttrs, RGBA, UINT, 8, PL);
        }
    } else if (args->imagetype == IMAGE_TYPE_RAW) {
        if(!args->pitchLinearOutput){
            NVM_SURF_FMT_SET_ATTR_RAW(surfFormatAttrs, GRBG, INT, 12, BL);
        } else {
            NVM_SURF_FMT_SET_ATTR_RAW(surfFormatAttrs, GRBG, INT, 12, PL);
        }
    } else{
        if(!args->pitchLinearOutput){
            NVM_SURF_FMT_SET_ATTR_YUV(surfFormatAttrs, YUV, 420, SEMI_PLANAR, UINT, 8, BL);
        } else {
            NVM_SURF_FMT_SET_ATTR_YUV(surfFormatAttrs, YUV, 420, SEMI_PLANAR, UINT, 8, PL);
        }
    }

    /* Get NvMediaSurfaceType for the specified surface format attributes */
    imgProducer->surfaceType = NvMediaSurfaceFormatGetType(surfFormatAttrs,  // surf format attrs
                                                           NVM_SURF_FMT_ATTR_MAX);               // num attrs

    LOG_DBG("%s: Surface type %d\n", __func__, imgProducer->surfaceType);

    imgProducer->device = NvMediaDeviceCreate();
    if(!imgProducer->device) {
        LOG_ERR("%s: Failed to create NvMedia device\n", __func__);
        return 1;
    }

    if(imgProducer->useblitpath) {
        imgProducer->blitter = NvMedia2DCreate(imgProducer->device);
        if(!imgProducer->blitter) {
            LOG_ERR("%s: Failed to create NvMedia 2D blitter\n", __func__);
            return 1;
        }
    }

    /* get the total number of frames*/
    if(!imgProducer->frameCount) {
        switch(args->imagetype){
            case IMAGE_TYPE_RGBA :
                imageSize = imgProducer->width * imgProducer->height * 4;
                break;
            case IMAGE_TYPE_RAW :
                imageSize = imgProducer->width * imgProducer->height * 2;
                break;
            case IMAGE_TYPE_YUV420 :
                imageSize = imgProducer->width * imgProducer->height * 1.5;
                break;
            default :
                LOG_ERR("Invalid image type\n");
                return 1;
        }

       imgProducer->frameCount = GetFramesCount(imgProducer->inpFileName,imageSize);
    }
    LOG_DBG("\n Frame count = %d\n",imgProducer->frameCount);

    /* create Producer Input Queue*/
    if(IsFailed(NvQueueCreate(&imgProducer->inputQueue,
                              IMAGE_BUFFERS_POOL_SIZE,
                              sizeof(NvMediaImage *)))) {
        return 1;
    }

    /* Set surface allocation attributes */
    surfAllocAttrs[0].type = NVM_SURF_ATTR_WIDTH;
    surfAllocAttrs[0].value = imgProducer->width;
    surfAllocAttrs[1].type = NVM_SURF_ATTR_HEIGHT;
    surfAllocAttrs[1].value = imgProducer->height;
    surfAllocAttrs[2].type = NVM_SURF_ATTR_EMB_LINES_TOP;
    surfAllocAttrs[2].value = 0;
    surfAllocAttrs[3].type = NVM_SURF_ATTR_EMB_LINES_BOTTOM;
    surfAllocAttrs[3].value = 0;
    numSurfAllocAttrs = 4;
    if ((env = getenv("DISPLAY_VM"))) {
        surfAllocAttrs[4].type = NVM_SURF_ATTR_PEER_VM_ID;
        surfAllocAttrs[4].value = atoi(env);
        numSurfAllocAttrs += 1;
    }

    for(i = 0; i < IMAGE_BUFFERS_POOL_SIZE; i++) {
        /* Call NvMediaImageCreateNew API for surface creation */
        image =  NvMediaImageCreateNew(imgProducer->device,             // device
                                      imgProducer->surfaceType,               // NvMediaSurfaceType type
                                      surfAllocAttrs,     // surf allocation attrs
                                      numSurfAllocAttrs,  // num attrs
                                      0);                 // flags
        if(!image) {
            return 1;
        }
        image->tag = imgProducer->inputQueue;

        if(IsFailed(NvQueuePut(imgProducer->inputQueue,
                (void *)&image, NV_TIMEOUT_INFINITE)))
            return 1;
    }

    //Create a NvMediaImage if Blit path has to be used
    if(imgProducer->useblitpath) {
        imgProducer->fileimage = NvMediaImageCreateNew(imgProducer->device,             // device
                                      imgProducer->surfaceType,               // NvMediaSurfaceType type
                                      surfAllocAttrs,     // surf allocation attrs
                                      numSurfAllocAttrs,                 // num attrs
                                      0);                 // flags
        if(!imgProducer->fileimage) {
            LOG_ERR("Failed in creating NvMediaImage for File Read\n");
            return 1;
        }
    }
    // Create EGLStream-Producer
    imgProducer->producer = NvMediaEglStreamProducerCreate(imgProducer->device,
                                                           imgProducer->display,
                                                           imgProducer->eglStream,
                                                           imgProducer->surfaceType,
                                                           imgProducer->width,
                                                           imgProducer->height);
    if(!imgProducer->producer) {
        LOG_ERR("ImageProducerInit: Unable to create producer\n");
        return 1;
    }

    pthread_create(&imgProducer->procThread, NULL,
                   ProcessImageThread, (void *)imgProducer);
    if(!imgProducer->procThread) {
        LOG_ERR("ImageProducerInit: Unable to create ProcessImageThread\n");
        return 1;
    }

    return 0;
}


void ImageProducerFini(ImageProducerCtx *imgProducer)
{
    NvMediaImage *image = NULL;
    LOG_DBG("ImageProducerFini: start\n");
    /*Destroy the input queue*/
    if(imgProducer->inputQueue) {
        while(IsSucceed(NvQueueGet(imgProducer->inputQueue, &image, 0))) {
            if (image) {
                NvMediaImageDestroy(image);
                image = NULL;
            }
        }
        LOG_DBG("\n Destroying producer input queue");
        NvQueueDestroy(imgProducer->inputQueue);
    }

    if(imgProducer->useblitpath && imgProducer->fileimage) {
        NvMediaImageDestroy(imgProducer->fileimage);
        imgProducer->fileimage = NULL;
    }

    if(imgProducer->useblitpath && imgProducer->blitter) {
        NvMedia2DDestroy(imgProducer->blitter);
        imgProducer->blitter = NULL;
    }

    if(imgProducer->producer) {
       NvMediaEglStreamProducerDestroy(imgProducer->producer);
       imgProducer->producer = NULL;
    }

    if(imgProducer->device) {
       NvMediaDeviceDestroy(imgProducer->device);
       imgProducer->device = NULL;
    }

    LOG_DBG("ImageProducerFini: end\n");
}

//Stop the imge producer thread
void ImageProducerStop(ImageProducerCtx *imgProducer)
{
    LOG_DBG("ImageProducerStop, wait for thread stop\n");
    if(imgProducer->procThread) {
        pthread_join(imgProducer->procThread, NULL);
    }
    LOG_DBG("ImageProducerStop, thread stop\n");
}

//Image producer to get the image back from EGLStream, and release them
void ImageProducerFlush(ImageProducerCtx *imgProducer)
{
    NvMediaImage *releaseImage = NULL;

    while (NvMediaEglStreamProducerGetImage(
           imgProducer->producer, &releaseImage, 0) == NVMEDIA_STATUS_OK) {
        NvQueuePut((NvQueue *)releaseImage->tag,
                  (void *)&releaseImage,
                  ENQUEUE_TIMEOUT);
    }
}
