/* 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 EGL stream sample app
//
//

// #define EGL_EGLEXT_PROTOTYPES

#include "cudaEGL.h"
#include "cuda_consumer.h"
#include "cuda_producer.h"
#include "eglstrm_common.h"

/* ------  globals ---------*/

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

#define NUM_TRAILS 4

bool signal_stop = 0;

static void sig_handler(int sig)
{
    signal_stop = 1;
    printf("Signal: %d\n", sig);
}

int main(int argc, char **argv)
{
    TestArgs     args;
    CUresult     curesult = CUDA_SUCCESS;
    unsigned int i, j;
    EGLint       streamState = 0;

    test_cuda_consumer_s cudaConsumer;
    test_cuda_producer_s cudaProducer;

    memset(&cudaProducer, 0, sizeof(test_cuda_producer_s));
    memset(&cudaConsumer, 0, sizeof(test_cuda_consumer_s));

    // Hook up Ctrl-C handler
    signal(SIGINT, sig_handler);
    if (!eglSetupExtensions()) {
        printf("SetupExtentions failed \n");
        curesult = CUDA_ERROR_UNKNOWN;
        goto done;
    }

    checkCudaErrors(cuInit(0));

    int count;

    checkCudaErrors(cuDeviceGetCount(&count));
    printf("Found %d cuda devices\n", count);

    CUdevice devId;

    if (!EGLStreamInit(&devId)) {
        printf("EGLStream Init failed.\n");
        curesult = CUDA_ERROR_UNKNOWN;
        goto done;
    }
    curesult = cudaDeviceCreateProducer(&cudaProducer, devId);
    if (curesult != CUDA_SUCCESS) {
        goto done;
    }
    curesult = cudaDeviceCreateConsumer(&cudaConsumer, devId);
    if (curesult != CUDA_SUCCESS) {
        goto done;
    }
    checkCudaErrors(cuCtxPushCurrent(cudaConsumer.context));
    if (CUDA_SUCCESS != (curesult = cuEGLStreamConsumerConnect(&(cudaConsumer.cudaConn), eglStream))) {
        printf("FAILED Connect CUDA consumer  with error %d\n", curesult);
        goto done;
    }
    else {
        printf("Connected CUDA consumer, CudaConsumer %p\n", cudaConsumer.cudaConn);
    }
    checkCudaErrors(cuCtxPopCurrent(&cudaConsumer.context));

    checkCudaErrors(cuCtxPushCurrent(cudaProducer.context));
    if (CUDA_SUCCESS == (curesult = cuEGLStreamProducerConnect(&(cudaProducer.cudaConn), eglStream, WIDTH, HEIGHT))) {
        printf("Connect CUDA producer Done, CudaProducer %p\n", cudaProducer.cudaConn);
    }
    else {
        printf("Connect CUDA producer FAILED with error %d\n", curesult);
        goto done;
    }
    checkCudaErrors(cuCtxPopCurrent(&cudaProducer.context));

    // Initialize producer
    for (i = 0; i < NUM_TRAILS; i++) {
        if (streamState != EGL_STREAM_STATE_CONNECTING_KHR) {
            if (!eglQueryStreamKHR(g_display, eglStream, EGL_STREAM_STATE_KHR, &streamState)) {
                printf("main: eglQueryStreamKHR EGL_STREAM_STATE_KHR failed\n");
                curesult = CUDA_ERROR_UNKNOWN;
                goto done;
            }
        }
        args.inputWidth  = WIDTH;
        args.inputHeight = HEIGHT;
        if (i % 2 != 0) {
            args.isARGB  = 1;
            args.infile1 = sdkFindFilePath("cuda_f_1.yuv", argv[0]);
            args.infile2 = sdkFindFilePath("cuda_f_2.yuv", argv[0]);
        }
        else {
            args.isARGB  = 0;
            args.infile1 = sdkFindFilePath("cuda_yuv_f_1.yuv", argv[0]);
            args.infile2 = sdkFindFilePath("cuda_yuv_f_2.yuv", argv[0]);
        }
        if ((i % 4) < 2) {
            args.pitchLinearOutput = 1;
        }
        else {
            args.pitchLinearOutput = 0;
        }

        checkCudaErrors(cuCtxPushCurrent(cudaProducer.context));
        cudaProducerInit(&cudaProducer, g_display, eglStream, &args);
        checkCudaErrors(cuCtxPopCurrent(&cudaProducer.context));

        checkCudaErrors(cuCtxPushCurrent(cudaConsumer.context));
        cuda_consumer_init(&cudaConsumer, &args);
        checkCudaErrors(cuCtxPopCurrent(&cudaConsumer.context));

        printf("main - Cuda Producer and Consumer Initialized.\n");

        for (j = 0; j < 2; j++) {
            printf("Running for %s frame and %s input\n",
                   args.isARGB ? "ARGB" : "YUV",
                   args.pitchLinearOutput ? "Pitchlinear" : "BlockLinear");
            if (j == 0) {
                checkCudaErrors(cuCtxPushCurrent(cudaProducer.context));
                curesult = cudaProducerTest(&cudaProducer, cudaProducer.fileName1);
                if (curesult != CUDA_SUCCESS) {
                    printf("Cuda Producer Test failed for frame = %d\n", j + 1);
                    goto done;
                }
                checkCudaErrors(cuCtxPopCurrent(&cudaProducer.context));
                checkCudaErrors(cuCtxPushCurrent(cudaConsumer.context));
                curesult = cudaConsumerTest(&cudaConsumer, cudaConsumer.outFile1);
                if (curesult != CUDA_SUCCESS) {
                    printf("Cuda Consumer Test failed for frame = %d\n", j + 1);
                    goto done;
                }
                checkCudaErrors(cuCtxPopCurrent(&cudaConsumer.context));
            }
            else {
                checkCudaErrors(cuCtxPushCurrent(cudaProducer.context));
                curesult = cudaProducerTest(&cudaProducer, cudaProducer.fileName2);
                if (curesult != CUDA_SUCCESS) {
                    printf("Cuda Producer Test failed for frame = %d\n", j + 1);
                    goto done;
                }

                checkCudaErrors(cuCtxPopCurrent(&cudaProducer.context));
                checkCudaErrors(cuCtxPushCurrent(cudaConsumer.context));
                curesult = cudaConsumerTest(&cudaConsumer, cudaConsumer.outFile2);
                if (curesult != CUDA_SUCCESS) {
                    printf("Cuda Consumer Test failed for frame = %d\n", j + 1);
                    goto done;
                }
                checkCudaErrors(cuCtxPopCurrent(&cudaConsumer.context));
            }
        }
    }

    checkCudaErrors(cuCtxPushCurrent(cudaProducer.context));
    if (CUDA_SUCCESS != (curesult = cudaProducerDeinit(&cudaProducer))) {
        printf("Producer Disconnect FAILED. \n");
        goto done;
    }
    checkCudaErrors(cuCtxPopCurrent(&cudaProducer.context));

    if (!eglQueryStreamKHR(g_display, eglStream, EGL_STREAM_STATE_KHR, &streamState)) {
        printf("Cuda consumer, eglQueryStreamKHR EGL_STREAM_STATE_KHR failed\n");
        curesult = CUDA_ERROR_UNKNOWN;
        goto done;
    }
    if (streamState != EGL_STREAM_STATE_DISCONNECTED_KHR) {
        if (CUDA_SUCCESS != (curesult = cuda_consumer_deinit(&cudaConsumer))) {
            printf("Consumer Disconnect FAILED.\n");
            goto done;
        }
    }
    printf("Producer and Consumer Disconnected \n");

done:
    if (!eglQueryStreamKHR(g_display, eglStream, EGL_STREAM_STATE_KHR, &streamState)) {
        printf("Cuda consumer, eglQueryStreamKHR EGL_STREAM_STATE_KHR failed\n");
        curesult = CUDA_ERROR_UNKNOWN;
    }
    if (streamState != EGL_STREAM_STATE_DISCONNECTED_KHR) {
        EGLStreamFini();
    }

    if (curesult == CUDA_SUCCESS) {
        printf("&&&& EGLStream interop test PASSED\n");
    }
    else {
        printf("&&&& EGLStream interop test FAILED\n");
    }
    return 0;
}
