/* 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.
 */

//
// DESCRIPTION:   Simple CUDA consumer rendering sample app
//

#include "cuda_consumer.h"

#include <helper_cuda_drvapi.h>

#include "eglstrm_common.h"

#if defined(EXTENSION_LIST)
EXTENSION_LIST(EXTLST_EXTERN)
#endif

int checkbuf(FILE *fp1, FILE *fp2);

CUresult cudaConsumerTest(test_cuda_consumer_s *data, const char *fileName)
{
    CUresult           cuStatus = CUDA_SUCCESS;
    CUarray            cudaArr  = NULL;
    CUeglFrame         cudaEgl;
    CUgraphicsResource cudaResource;
    unsigned int       i;
    int                check_result;
    FILE              *pInFile1 = NULL, *pInFile2 = NULL, *file_p = NULL;
    EGLint             streamState = 0;

    if (!data) {
        printf("%s: Bad parameter\n", __func__);
        goto done;
    }

    if (!eglQueryStreamKHR(g_display, eglStream, EGL_STREAM_STATE_KHR, &streamState)) {
        printf("Cuda consumer, eglQueryStreamKHR EGL_STREAM_STATE_KHR failed\n");
    }
    if (streamState == EGL_STREAM_STATE_DISCONNECTED_KHR) {
        printf("CUDA Consumer: - EGL_STREAM_STATE_DISCONNECTED_KHR received\n");
    }

    if (streamState == EGL_STREAM_STATE_NEW_FRAME_AVAILABLE_KHR) {
        cuStatus = cuEGLStreamConsumerAcquireFrame(&(data->cudaConn), &cudaResource, NULL, 16000);

        if (cuStatus == CUDA_SUCCESS) {
            CUdeviceptr    pDevPtr = 0;
            int            bufferSize;
            unsigned char *pCudaCopyMem     = NULL;
            unsigned int   copyWidthInBytes = 0, copyHeight = 0;

            file_p = fopen(fileName, "wb+");
            if (!file_p) {
                printf("WriteFrame: file open failed %s\n", fileName);
                cuStatus = CUDA_ERROR_UNKNOWN;
                goto done;
            }
            cuStatus = cuGraphicsResourceGetMappedEglFrame(&cudaEgl, cudaResource, 0, 0);
            if (cuStatus != CUDA_SUCCESS) {
                printf("Cuda get resource failed with %d\n", cuStatus);
                goto done;
            }
            cuStatus = cuCtxSynchronize();
            if (cuStatus != CUDA_SUCCESS) {
                printf("cuCtxSynchronize failed \n");
                goto done;
            }
            if (!(cudaEgl.planeCount >= 1 && cudaEgl.planeCount <= 3)) {
                printf("Plane count is invalid\nExiting\n");
                goto done;
            }

            for (i = 0; i < cudaEgl.planeCount; i++) {
                if (cudaEgl.frameType == CU_EGL_FRAME_TYPE_PITCH) {
                    pDevPtr = (CUdeviceptr)cudaEgl.frame.pPitch[i];
                    if (cudaEgl.planeCount == 1) {
                        bufferSize       = cudaEgl.pitch * cudaEgl.height;
                        copyWidthInBytes = cudaEgl.pitch;
                        copyHeight       = data->height;
                    }
                    else if (i == 1 && cudaEgl.planeCount == 2) { // YUV 420
                                                                  // semi-planar
                        bufferSize       = cudaEgl.pitch * cudaEgl.height / 2;
                        copyWidthInBytes = cudaEgl.pitch;
                        copyHeight       = data->height / 2;
                    }
                    else {
                        bufferSize       = data->width * data->height;
                        copyWidthInBytes = data->width;
                        copyHeight       = data->height;
                        if (i > 0) {
                            bufferSize >>= 2;
                            copyWidthInBytes >>= 1;
                            copyHeight >>= 1;
                        }
                    }
                }
                else {
                    cudaArr = cudaEgl.frame.pArray[i];
                    if (cudaEgl.planeCount == 1) {
                        bufferSize       = data->width * data->height * 4;
                        copyWidthInBytes = data->width * 4;
                        copyHeight       = data->height;
                    }
                    else if (i == 1 && cudaEgl.planeCount == 2) { // YUV 420
                                                                  // semi-planar
                        bufferSize       = data->width * data->height / 2;
                        copyWidthInBytes = data->width;
                        copyHeight       = data->height / 2;
                    }
                    else {
                        bufferSize       = data->width * data->height;
                        copyWidthInBytes = data->width;
                        copyHeight       = data->height;
                        if (i > 0) {
                            bufferSize >>= 2;
                            copyWidthInBytes >>= 1;
                            copyHeight >>= 1;
                        }
                    }
                }
                if (i == 0) {
                    pCudaCopyMem = (unsigned char *)malloc(bufferSize);
                    if (pCudaCopyMem == NULL) {
                        printf("pCudaCopyMem malloc failed\n");
                        goto done;
                    }
                }
                memset(pCudaCopyMem, 0, bufferSize);
                if (data->pitchLinearOutput) {
                    cuStatus = cuMemcpyDtoH(pCudaCopyMem, pDevPtr, bufferSize);
                    if (cuStatus != CUDA_SUCCESS) {
                        printf("cuda_consumer: pitch linear Memcpy failed, bufferSize =%d\n", bufferSize);
                        goto done;
                    }
                    cuStatus = cuCtxSynchronize();
                    if (cuStatus != CUDA_SUCCESS) {
                        printf("cuda_consumer: cuCtxSynchronize failed after memcpy \n");
                        goto done;
                    }
                }
                else {
                    CUDA_MEMCPY3D cpdesc;
                    memset(&cpdesc, 0, sizeof(cpdesc));
                    cpdesc.srcXInBytes = cpdesc.srcY = cpdesc.srcZ = cpdesc.srcLOD = 0;
                    cpdesc.srcMemoryType                                           = CU_MEMORYTYPE_ARRAY;
                    cpdesc.srcArray                                                = cudaArr;
                    cpdesc.dstXInBytes = cpdesc.dstY = cpdesc.dstZ = cpdesc.dstLOD = 0;
                    cpdesc.dstMemoryType                                           = CU_MEMORYTYPE_HOST;
                    cpdesc.dstHost                                                 = (void *)pCudaCopyMem;
                    cpdesc.WidthInBytes = copyWidthInBytes; // data->width * 4;
                    cpdesc.Height       = copyHeight;       // data->height;
                    cpdesc.Depth        = 1;

                    cuStatus = cuMemcpy3D(&cpdesc);
                    if (cuStatus != CUDA_SUCCESS) {
                        printf("Cuda consumer: cuMemCpy3D failed,  copyWidthInBytes=%d, "
                               "copyHight=%d\n",
                               copyWidthInBytes,
                               copyHeight);
                    }
                    cuStatus = cuCtxSynchronize();
                    if (cuStatus != CUDA_SUCCESS) {
                        printf("cuCtxSynchronize failed after memcpy \n");
                    }
                }
                if (cuStatus == CUDA_SUCCESS) {
                    if (fwrite(pCudaCopyMem, bufferSize, 1, file_p) != 1) {
                        printf("Cuda consumer: output file write failed\n");
                        cuStatus = CUDA_ERROR_UNKNOWN;
                        goto done;
                    }
                }
            }
            pInFile1 = fopen(data->fileName1, "rb");
            if (!pInFile1) {
                printf("Failed to open file :%s\n", data->fileName1);
                goto done;
            }
            pInFile2 = fopen(data->fileName2, "rb");
            if (!pInFile2) {
                printf("Failed to open file :%s\n", data->fileName2);
                goto done;
            }
            rewind(file_p);
            check_result = checkbuf(file_p, pInFile1);
            if (check_result == -1) {
                rewind(file_p);
                check_result = checkbuf(file_p, pInFile2);
                if (check_result == -1) {
                    printf("Frame received does not match any valid image: FAILED\n");
                }
                else {
                    printf("Frame check Passed\n");
                }
            }
            else {
                printf("Frame check Passed\n");
            }
            if (pCudaCopyMem) {
                free(pCudaCopyMem);
                pCudaCopyMem = NULL;
            }
            cuStatus = cuEGLStreamConsumerReleaseFrame(&data->cudaConn, cudaResource, NULL);
            if (cuStatus != CUDA_SUCCESS) {
                printf("cuEGLStreamConsumerReleaseFrame failed with cuStatus = %d\n", cuStatus);
                goto done;
            }
        }
        else {
            printf("cuda AcquireFrame FAILED with  cuStatus=%d\n", cuStatus);
            goto done;
        }
    }

done:
    if (file_p) {
        fclose(file_p);
        file_p = NULL;
    }
    if (pInFile1) {
        fclose(pInFile1);
        pInFile1 = NULL;
    }
    if (pInFile1) {
        fclose(pInFile2);
        pInFile2 = NULL;
    }
    return cuStatus;
}

int checkbuf(FILE *fp1, FILE *fp2)
{
    int match = 0;
    int ch1, ch2;
    if (fp1 == NULL) {
        printf("Invalid file pointer for first file\n");
        return -1;
    }
    else if (fp2 == NULL) {
        printf("Invalid file pointer for second file\n");
        return -1;
    }
    else {
        ch1 = getc(fp1);
        ch2 = getc(fp2);
        while ((ch1 != EOF) && (ch2 != EOF) && (ch1 == ch2)) {
            ch1 = getc(fp1);
            ch2 = getc(fp2);
        }
        if (ch1 == ch2) {
            match = 1;
        }
        else if (ch1 != ch2) {
            match = -1;
        }
    }
    return match;
}

CUresult cudaDeviceCreateConsumer(test_cuda_consumer_s *cudaConsumer, CUdevice device)
{
    CUresult status = CUDA_SUCCESS;
    if (CUDA_SUCCESS != (status = cuInit(0))) {
        printf("Failed to initialize CUDA\n");
        return status;
    }

    int  major = 0, minor = 0;
    char deviceName[256];
    checkCudaErrors(cuDeviceGetAttribute(&major, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, device));
    checkCudaErrors(cuDeviceGetAttribute(&minor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, device));
    checkCudaErrors(cuDeviceGetName(deviceName, 256, device));
    printf("CUDA Consumer on GPU Device %d: \"%s\" with compute capability "
           "%d.%d\n\n",
           device,
           deviceName,
           major,
           minor);

    CUctxCreateParams ctxCreateParams = {};
    if (CUDA_SUCCESS != (status = cuCtxCreate(&cudaConsumer->context, &ctxCreateParams, 0, device))) {
        printf("failed to create CUDA context\n");
        return status;
    }
    checkCudaErrors(cuCtxPopCurrent(&cudaConsumer->context));
    return status;
}

void cuda_consumer_init(test_cuda_consumer_s *cudaConsumer, TestArgs *args)
{
    cudaConsumer->pitchLinearOutput = args->pitchLinearOutput;
    cudaConsumer->width             = args->inputWidth;
    cudaConsumer->height            = args->inputHeight;
    cudaConsumer->fileName1         = args->infile1;
    cudaConsumer->fileName2         = args->infile2;

    cudaConsumer->outFile1 = const_cast<char *>("cuda_out1.yuv");
    cudaConsumer->outFile2 = const_cast<char *>("cuda_out2.yuv");
}

CUresult cuda_consumer_deinit(test_cuda_consumer_s *cudaConsumer)
{
    return cuEGLStreamConsumerDisconnect(&cudaConsumer->cudaConn);
}
