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

#include "cudaEGL.h"
#include "cuda_consumer.h"
#include "cuda_producer.h"
#include "eglstrm_common.h"
#include "helper.h"
#if defined(EXTENSION_LIST)
EXTENSION_LIST(EXTLST_EXTERN)
#endif

bool        signal_stop = 0;
extern bool verbose;

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

void DoneCons(int consumerStatus, int send_fd)
{
    EGLStreamFini();
    // get the final status from producer, combine and print
    int producerStatus = -1;
    if (-1 == recv(send_fd, (void *)&producerStatus, sizeof(int), 0)) {
        printf("%s: Cuda Consumer could not receive status from producer.\n", __func__);
    }
    close(send_fd);

    if (producerStatus == 0 && consumerStatus == 0) {
        printf("&&&& EGLStream_CUDA_CrossGPU PASSED\n");
        exit(EXIT_SUCCESS);
    }
    else {
        printf("&&&& EGLStream_CUDA_CrossGPU FAILED\n");
        exit(EXIT_FAILURE);
    }
}

void DoneProd(int producerStatus, int connect_fd)
{
    EGLStreamFini();
    if (-1 == send(connect_fd, (void *)&producerStatus, sizeof(int), 0)) {
        printf("%s: Cuda Producer could not send status to consumer.\n", __func__);
    }
    close(connect_fd);
    if (producerStatus == 0) {
        exit(EXIT_SUCCESS);
    }
    else {
        exit(EXIT_FAILURE);
    }
}

int WIDTH = 8192, HEIGHT = 8192;
int main(int argc, char **argv)
{
    TestArgs                   args           = {0, false};
    CUresult                   curesult       = CUDA_SUCCESS;
    unsigned int               j              = 0;
    cudaError_t                err            = cudaSuccess;
    EGLNativeFileDescriptorKHR fileDescriptor = EGL_NO_FILE_DESCRIPTOR_KHR;
    struct timespec            start, end;
    CUeglFrame                 cudaEgl1, cudaEgl2;
    int                        consumerStatus = 0;
    int                        send_fd        = -1;

    if (parseCmdLine(argc, argv, &args) < 0) {
        printUsage();
        curesult = CUDA_ERROR_UNKNOWN;
        DoneCons(consumerStatus, send_fd);
    }

    printf("Width : %u, height: %u and iterations: %u\n", WIDTH, HEIGHT, NUMTRIALS);

    if (!args.isProducer) // Consumer code
    {
        test_cuda_consumer_s cudaConsumer;
        memset(&cudaConsumer, 0, sizeof(test_cuda_consumer_s));
        cudaConsumer.profileAPI = profileAPIs;

        // Hook up Ctrl-C handler
        signal(SIGINT, sig_handler);

        if (!EGLStreamInit(isCrossDevice, !args.isProducer, EGL_NO_FILE_DESCRIPTOR_KHR)) {
            printf("EGLStream Init failed.\n");
            curesult       = CUDA_ERROR_UNKNOWN;
            consumerStatus = -1;
            DoneCons(consumerStatus, send_fd);
        }

        cudaConsumer.cudaDevId = cudaDevIndexCons;
        curesult               = cudaDeviceCreateConsumer(&cudaConsumer);
        if (curesult != CUDA_SUCCESS) {
            consumerStatus = -1;
            DoneCons(consumerStatus, send_fd);
        }

        cuCtxPushCurrent(cudaConsumer.context);

        launchProducer(&args);

        args.charCnt = WIDTH * HEIGHT * 4;

        curesult = cuda_consumer_init(&cudaConsumer, &args);
        if (curesult != CUDA_SUCCESS) {
            printf("Cuda Consumer: Init failed, status: %d\n", curesult);
            consumerStatus = -1;
            DoneCons(consumerStatus, send_fd);
        }

        cuCtxPopCurrent(&cudaConsumer.context);

        send_fd = UnixSocketConnect(SOCK_PATH);
        if (-1 == send_fd) {
            printf("%s: Cuda Consumer cannot create socket %s\n", __func__, SOCK_PATH);
            consumerStatus = -1;
            DoneCons(consumerStatus, send_fd);
        }

        cuCtxPushCurrent(cudaConsumer.context);
        cudaConsumer.eglStream  = g_consumerEglStream;
        cudaConsumer.eglDisplay = g_consumerEglDisplay;

        // Send the EGL stream FD to producer
        fileDescriptor = eglGetStreamFileDescriptorKHR(cudaConsumer.eglDisplay, cudaConsumer.eglStream);
        if (EGL_NO_FILE_DESCRIPTOR_KHR == fileDescriptor) {
            printf("%s: Cuda Consumer could not get EGL file descriptor.\n", __func__);
            eglDestroyStreamKHR(cudaConsumer.eglDisplay, cudaConsumer.eglStream);
            consumerStatus = -1;
            DoneCons(consumerStatus, send_fd);
        }

        if (verbose)
            printf("%s: Cuda Consumer EGL stream FD obtained : %d.\n", __func__, fileDescriptor);

        int res = -1;
        res     = EGLStreamSendfd(send_fd, fileDescriptor);
        if (-1 == res) {
            printf("%s: Cuda Consumer could not send EGL file descriptor.\n", __func__);
            consumerStatus = -1;
            close(fileDescriptor);
        }

        if (CUDA_SUCCESS != (curesult = cuEGLStreamConsumerConnect(&(cudaConsumer.cudaConn), cudaConsumer.eglStream))) {
            printf("FAILED Connect CUDA consumer with error %d\n", curesult);
            consumerStatus = -1;
            DoneCons(consumerStatus, send_fd);
        }

        j = 0;
        for (j = 0; j < NUMTRIALS; j++) {
            curesult = cudaConsumerAcquireFrame(&cudaConsumer, j);
            if (curesult != CUDA_SUCCESS) {
                printf("Cuda Consumer Test failed for frame = %d\n", j + 1);
                consumerStatus = -1;
                DoneCons(consumerStatus, send_fd);
            }
            curesult = cudaConsumerReleaseFrame(&cudaConsumer, j);
            if (curesult != CUDA_SUCCESS) {
                printf("Cuda Consumer Test failed for frame = %d\n", j + 1);
                consumerStatus = -1;
                DoneCons(consumerStatus, send_fd);
            }

            curesult = cudaConsumerAcquireFrame(&cudaConsumer, j);
            if (curesult != CUDA_SUCCESS) {
                printf("Cuda Consumer Test failed for frame = %d\n", j + 1);
                consumerStatus = -1;
                DoneCons(consumerStatus, send_fd);
            }
            curesult = cudaConsumerReleaseFrame(&cudaConsumer, j);
            if (curesult != CUDA_SUCCESS) {
                printf("Cuda Consumer Test failed for frame = %d\n", j + 1);
                consumerStatus = -1;
                DoneCons(consumerStatus, send_fd);
            }
        }
        cuCtxSynchronize();
        close(fileDescriptor);
        err = cudaGetValueMismatch();
        if (err != cudaSuccess) {
            printf("Consumer: App failed with value mismatch\n");
            curesult       = CUDA_ERROR_UNKNOWN;
            consumerStatus = -1;
            DoneCons(consumerStatus, send_fd);
        }

        EGLint streamState = 0;
        if (!eglQueryStreamKHR(cudaConsumer.eglDisplay, cudaConsumer.eglStream, EGL_STREAM_STATE_KHR, &streamState)) {
            printf("Main, eglQueryStreamKHR EGL_STREAM_STATE_KHR failed\n");
            curesult       = CUDA_ERROR_UNKNOWN;
            consumerStatus = -1;
            DoneCons(consumerStatus, send_fd);
        }

        if (streamState != EGL_STREAM_STATE_DISCONNECTED_KHR) {
            if (CUDA_SUCCESS != (curesult = cuda_consumer_Deinit(&cudaConsumer))) {
                printf("Consumer Disconnect FAILED.\n");
                consumerStatus = -1;
                DoneCons(consumerStatus, send_fd);
            }
        }
    }
    else // Producer
    {
        test_cuda_producer_s cudaProducer;
        memset(&cudaProducer, 0, sizeof(test_cuda_producer_s));
        cudaProducer.profileAPI = profileAPIs;
        int producerStatus      = 0;

        setenv("CUDA_EGL_PRODUCER_RETURN_WAIT_TIMEOUT", "1600", 0);

        int connect_fd = -1;
        // Hook up Ctrl-C handler
        signal(SIGINT, sig_handler);

        // Create connection to Consumer
        connect_fd = UnixSocketCreate(SOCK_PATH);
        if (-1 == connect_fd) {
            printf("%s: Cuda Producer could not create socket: %s.\n", __func__, SOCK_PATH);
            producerStatus = -1;
            DoneProd(producerStatus, connect_fd);
        }

        // Get the file descriptor of the stream from the consumer process
        // and re-create the EGL stream from it
        fileDescriptor = EGLStreamReceivefd(connect_fd);
        if (-1 == fileDescriptor) {
            printf("%s: Cuda Producer could not receive EGL file descriptor \n", __func__);
            producerStatus = -1;
            DoneProd(producerStatus, connect_fd);
        }

        if (!EGLStreamInit(isCrossDevice, 0, fileDescriptor)) {
            printf("EGLStream Init failed.\n");
            producerStatus = -1;
            curesult       = CUDA_ERROR_UNKNOWN;
            DoneProd(producerStatus, connect_fd);
        }

        cudaProducer.eglDisplay = g_producerEglDisplay;
        cudaProducer.eglStream  = g_producerEglStream;
        cudaProducer.cudaDevId  = cudaDevIndexProd;

        curesult = cudaDeviceCreateProducer(&cudaProducer);
        if (curesult != CUDA_SUCCESS) {
            producerStatus = -1;
            DoneProd(producerStatus, connect_fd);
        }

        args.charCnt = WIDTH * HEIGHT * 4;
        cuCtxPushCurrent(cudaProducer.context);
        curesult = cudaProducerInit(&cudaProducer, &args);
        if (curesult != CUDA_SUCCESS) {
            printf("Cuda Producer: Init failed, status: %d\n", curesult);
            producerStatus = -1;
            DoneProd(producerStatus, connect_fd);
        }

        // wait for consumer to connect first
        int    err         = 0;
        int    wait_loop   = 0;
        EGLint streamState = 0;
        do {
            err =
                eglQueryStreamKHR(cudaProducer.eglDisplay, cudaProducer.eglStream, EGL_STREAM_STATE_KHR, &streamState);
            if ((0 != err) && (EGL_STREAM_STATE_CONNECTING_KHR != streamState)) {
                sleep(1);
                wait_loop++;
            }
        } while ((wait_loop < 10) && (0 != err) && (streamState != EGL_STREAM_STATE_CONNECTING_KHR));

        if ((0 == err) || (wait_loop >= 10)) {
            printf("%s: Cuda Producer eglQueryStreamKHR EGL_STREAM_STATE_KHR failed.\n", __func__);
            producerStatus = -1;
            DoneProd(producerStatus, connect_fd);
        }

        if (CUDA_SUCCESS
            != (curesult =
                    cuEGLStreamProducerConnect(&(cudaProducer.cudaConn), cudaProducer.eglStream, WIDTH, HEIGHT))) {
            printf("Connect CUDA producer FAILED with error %d\n", curesult);
            producerStatus = -1;
            DoneProd(producerStatus, connect_fd);
        }

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

        cudaProducerPrepareFrame(&cudaEgl1, cudaProducer.cudaPtr, args.charCnt);
        cudaProducerPrepareFrame(&cudaEgl2, cudaProducer.cudaPtr1, args.charCnt);

        j = 0;
        for (j = 0; j < NUMTRIALS; j++) {
            curesult = cudaProducerPresentFrame(&cudaProducer, cudaEgl1, j);
            if (curesult != CUDA_SUCCESS) {
                printf("Cuda Producer Test failed for frame = %d with cuda error:%d\n", j + 1, curesult);
                producerStatus = -1;
                DoneProd(producerStatus, connect_fd);
            }

            curesult = cudaProducerPresentFrame(&cudaProducer, cudaEgl2, j);
            if (curesult != CUDA_SUCCESS) {
                printf("Cuda Producer Test failed for frame = %d with cuda error:%d\n", j + 1, curesult);
                producerStatus = -1;
                DoneProd(producerStatus, connect_fd);
            }

            curesult = cudaProducerReturnFrame(&cudaProducer, cudaEgl1, j);
            if (curesult != CUDA_SUCCESS) {
                printf("Cuda Producer Test failed for frame = %d with cuda error:%d\n", j + 1, curesult);
                producerStatus = -1;
                DoneProd(producerStatus, connect_fd);
            }

            curesult = cudaProducerReturnFrame(&cudaProducer, cudaEgl2, j);
            if (curesult != CUDA_SUCCESS) {
                printf("Cuda Producer Test failed for frame = %d with cuda error:%d\n", j + 1, curesult);
                producerStatus = -1;
                DoneProd(producerStatus, connect_fd);
            }
        }

        cuCtxSynchronize();
        err = cudaGetValueMismatch();
        if (err != cudaSuccess) {
            printf("Prod: App failed with value mismatch\n");
            curesult       = CUDA_ERROR_UNKNOWN;
            producerStatus = -1;
            DoneProd(producerStatus, connect_fd);
        }

        printf("Tear Down Start.....\n");
        if (!eglQueryStreamKHR(cudaProducer.eglDisplay, cudaProducer.eglStream, EGL_STREAM_STATE_KHR, &streamState)) {
            printf("Main, eglQueryStreamKHR EGL_STREAM_STATE_KHR failed\n");
            curesult       = CUDA_ERROR_UNKNOWN;
            producerStatus = -1;
            DoneProd(producerStatus, connect_fd);
        }

        if (streamState != EGL_STREAM_STATE_DISCONNECTED_KHR) {
            if (CUDA_SUCCESS != (curesult = cudaProducerDeinit(&cudaProducer))) {
                printf("Producer Disconnect FAILED with %d\n", curesult);
                producerStatus = -1;
                DoneProd(producerStatus, connect_fd);
            }
        }
        unsetenv("CUDA_EGL_PRODUCER_RETURN_WAIT_TIMEOUT");
    }

    return 0;
}
