/* 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:   Common egl stream functions
//

#include "eglstrm_common.h"

EGLStreamKHR g_producerEglStream  = EGL_NO_STREAM_KHR;
EGLStreamKHR g_consumerEglStream  = EGL_NO_STREAM_KHR;
EGLDisplay   g_producerEglDisplay = EGL_NO_DISPLAY;
EGLDisplay   g_consumerEglDisplay = EGL_NO_DISPLAY;
int          cudaDevIndexProd     = -1;
int          cudaDevIndexCons     = -1;

#if defined(EXTENSION_LIST)
EXTENSION_LIST(EXTLST_DECL)
typedef void (*extlst_fnptr_t)(void);
static struct
{
    extlst_fnptr_t *fnptr;
    char const     *name;
    bool            is_dgpu; // This function is need only for dgpu case
} extensionList[] = {EXTENSION_LIST(EXTLST_ENTRY)};

int eglSetupExtensions(bool isCrossDevice)
{
    unsigned int i;

    for (i = 0; i < (sizeof(extensionList) / sizeof(*extensionList)); i++) {
        // load the dgpu function only if we are running cross device test
        if ((!extensionList[i].is_dgpu) || (extensionList[i].is_dgpu == isCrossDevice)) {
            *extensionList[i].fnptr = eglGetProcAddress(extensionList[i].name);
            if (*extensionList[i].fnptr == NULL) {
                printf("Couldn't get address of %s()\n", extensionList[i].name);
                return 0;
            }
        }
    }

    return 1;
}

int EGLStreamInit(bool isCrossDevice, int isConsumer, EGLNativeFileDescriptorKHR fileDesc)
{
    static const EGLint streamAttrFIFOMode[] = {
        EGL_STREAM_FIFO_LENGTH_KHR, 5, EGL_SUPPORT_REUSE_NV, EGL_FALSE, EGL_NONE};
    EGLDisplay   eglDisplay[2] = {0};
    EGLStreamKHR eglStream[2]  = {0};
    EGLBoolean   eglStatus;

#define MAX_EGL_DEVICES 4

    EGLDeviceEXT devices[MAX_EGL_DEVICES];
    EGLint       numDevices = 0;

    eglStatus = eglQueryDevicesEXT(MAX_EGL_DEVICES, devices, &numDevices);
    if (eglStatus != EGL_TRUE) {
        printf("Error querying EGL devices\n");
        goto Done;
    }

    if (numDevices == 0) {
        printf("No EGL devices found\n");
        eglStatus = EGL_FALSE;
        goto Done;
    }

    // If cross device, create discrete GPU stream first and then create the
    // integrated GPU stream to connect to it via fd. The other way round fails
    // in producer connect.
    //
    // TODO: Find out if this EGL behavior is by design.
    if (isConsumer) {
        int egl_device_id = 0;
        for (egl_device_id = 0; egl_device_id < numDevices; egl_device_id++) {
            EGLAttrib cuda_device;
            eglStatus = eglQueryDeviceAttribEXT(devices[egl_device_id], EGL_CUDA_DEVICE_NV, &cuda_device);
            if (eglStatus == EGL_TRUE) {
                cudaDevIndexCons = cuda_device; // We select first EGL-CUDA Capable
                                                // device for consumer.
                printf("Found EGL-CUDA Capable device with CUDA Device id = %d out of "
                       "egl_device_id = %d\n",
                       (int)cudaDevIndexCons,
                       egl_device_id);
                break;
            }
        }

        if (egl_device_id >= numDevices) {
            printf("No CUDA Capable EGL Device found.. Waiving execution\n");
            goto Done;
        }

        g_consumerEglDisplay = eglGetPlatformDisplayEXT(EGL_PLATFORM_DEVICE_EXT, (void *)devices[egl_device_id], NULL);
        if (g_consumerEglDisplay == EGL_NO_DISPLAY) {
            printf("Could not get EGL display from device. \n");
            eglStatus = EGL_FALSE;
            goto Done;
        }

        eglStatus = eglInitialize(g_consumerEglDisplay, 0, 0);
        if (!eglStatus) {
            printf("EGL failed to initialize. \n");
            eglStatus = EGL_FALSE;
            goto Done;
        }

        g_consumerEglStream = eglCreateStreamKHR(g_consumerEglDisplay, streamAttrFIFOMode);
        if (g_consumerEglStream == EGL_NO_STREAM_KHR) {
            printf("Could not create EGL stream.\n");
            eglStatus = EGL_FALSE;
            goto Done;
        }

        eglStatus = eglStreamAttribKHR(g_consumerEglDisplay, g_consumerEglStream, EGL_CONSUMER_LATENCY_USEC_KHR, 16000);
        if (eglStatus != EGL_TRUE) {
            printf("eglStreamAttribKHR EGL_CONSUMER_LATENCY_USEC_KHR failed\n");
            goto Done;
        }

        eglStatus =
            eglStreamAttribKHR(g_consumerEglDisplay, g_consumerEglStream, EGL_CONSUMER_ACQUIRE_TIMEOUT_USEC_KHR, 16000);
        if (eglStatus != EGL_TRUE) {
            printf("eglStreamAttribKHR EGL_CONSUMER_ACQUIRE_TIMEOUT_USEC_KHR failed\n");
            goto Done;
        }
    }

    if (!isConsumer) { // Producer

        if (fileDesc == EGL_NO_FILE_DESCRIPTOR_KHR) {
            printf("Cuda Producer received bad file descriptor\n");
            eglStatus = EGL_FALSE;
            goto Done;
        }

        int egl_device_id    = 0;
        int egl_cuda_devices = 0;
        for (egl_device_id = 0; egl_device_id < numDevices; egl_device_id++) {
            EGLAttrib cuda_device = -1;
            eglStatus             = eglQueryDeviceAttribEXT(devices[egl_device_id], EGL_CUDA_DEVICE_NV, &cuda_device);
            if (eglStatus == EGL_TRUE) {
                egl_cuda_devices++;
                if (isCrossDevice && (egl_cuda_devices > 1)) {
                    // We select second EGL-CUDA Capable device for producer.
                    cudaDevIndexProd = (int)cuda_device;
                    printf("Found EGL-CUDA Capable device with CUDA Device id = %d "
                           "egl_device_id = %d \n",
                           (int)cudaDevIndexProd,
                           egl_device_id);
                    break;
                }
                if (!isCrossDevice) {
                    // We select first EGL-CUDA Capable device for producer same as
                    // consumer.
                    cudaDevIndexProd = (int)cuda_device;
                    printf("Found EGL-CUDA Capable device with CUDA Device id = %d "
                           "egl_device_id = %d \n",
                           (int)cudaDevIndexProd,
                           egl_device_id);
                    break;
                }
            }
        }

        if (egl_device_id >= numDevices) {
            printf("No CUDA Capable EGL Device found.. Waiving execution\n");
            goto Done;
        }

        g_producerEglDisplay = eglGetPlatformDisplayEXT(EGL_PLATFORM_DEVICE_EXT, (void *)devices[egl_device_id], NULL);
        if (g_producerEglDisplay == EGL_NO_DISPLAY) {
            printf("Could not get EGL display from device. \n");
            eglStatus = EGL_FALSE;
            goto Done;
        }

        eglStatus = eglInitialize(g_producerEglDisplay, 0, 0);
        if (!eglStatus) {
            printf("EGL failed to initialize. \n");
            eglStatus = EGL_FALSE;
            goto Done;
        }

        g_producerEglStream = eglCreateStreamFromFileDescriptorKHR(g_producerEglDisplay, fileDesc);
        close(fileDesc);

        if (g_producerEglStream == EGL_NO_STREAM_KHR) {
            printf("CUDA Producer Could not create EGL stream.\n");
            eglStatus = EGL_FALSE;
            goto Done;
        }
        else {
            printf("Producer created EGLStream for the GPU.\n");
        }
    }

Done:
    return eglStatus == EGL_TRUE ? 1 : 0;
}

void EGLStreamFini(void)
{
    if (g_producerEglStream != EGL_NO_STREAM_KHR) {
        eglDestroyStreamKHR(g_producerEglDisplay, g_producerEglStream);
    }
    if (g_consumerEglStream != g_producerEglStream) {
        if (g_consumerEglStream != EGL_NO_STREAM_KHR) {
            eglDestroyStreamKHR(g_consumerEglDisplay, g_consumerEglStream);
        }
    }
}

int UnixSocketConnect(const char *socket_name)
{
    int                sock_fd = -1;
    struct sockaddr_un sock_addr;
    int                wait_loop = 0;

    sock_fd = socket(PF_UNIX, SOCK_STREAM, 0);
    if (sock_fd < 0) {
        printf("%s: socket create failed.\n", __func__);
        return -1;
    }

    if (verbose)
        printf("%s: send_fd: sock_fd: %d\n", __func__, sock_fd);

    memset(&sock_addr, 0, sizeof(struct sockaddr_un));
    sock_addr.sun_family = AF_UNIX;
    strncpy(sock_addr.sun_path, socket_name, sizeof(sock_addr.sun_path) - 1);

    while (connect(sock_fd, (const struct sockaddr *)&sock_addr, sizeof(struct sockaddr_un))) {
        if (wait_loop < 60) {
            if (!wait_loop)
                printf("Waiting for EGL stream producer ");
            else
                printf(".");
            fflush(stdout);
            sleep(1);
            wait_loop++;
        }
        else {
            printf("\n%s: Waiting timed out\n", __func__);
            return -1;
        }
    }
    if (wait_loop)
        printf("\n");

    if (verbose)
        printf("%s: Wait is done\n", __func__);

    return sock_fd;
}

/* Send <fd_to_send> (a file descriptor) to another process */
/* over a unix domain socket named <socket_name>.           */
/* <socket_name> can be any nonexistant filename.           */
int EGLStreamSendfd(int send_fd, int fd_to_send)
{
    struct msghdr   msg;
    struct iovec    iov[1];
    char            ctrl_buf[CMSG_SPACE(sizeof(int))];
    struct cmsghdr *cmsg = NULL;
    void           *data;
    int             res;
    memset(&msg, 0, sizeof(msg));

    iov[0].iov_len  = 1;           // must send at least 1 byte
    iov[0].iov_base = (void *)"x"; // any byte value (value ignored)
    msg.msg_iov     = iov;
    msg.msg_iovlen  = 1;

    memset(ctrl_buf, 0, sizeof(ctrl_buf));
    msg.msg_control    = ctrl_buf;
    msg.msg_controllen = sizeof(ctrl_buf);

    cmsg             = CMSG_FIRSTHDR(&msg);
    cmsg->cmsg_level = SOL_SOCKET;
    cmsg->cmsg_type  = SCM_RIGHTS;
    cmsg->cmsg_len   = CMSG_LEN(sizeof(int));
    data             = CMSG_DATA(cmsg);
    *(int *)data     = fd_to_send;

    msg.msg_controllen = cmsg->cmsg_len;

    res = sendmsg(send_fd, &msg, 0);
    if (res <= 0) {
        printf("%s: sendmsg failed", __func__);
        return -1;
    }

    return 0;
}

/* Listen on a unix domain socket named <socket_name>.     */
/* Connect to it and return connect_fd                     */
int UnixSocketCreate(const char *socket_name)
{
    int                listen_fd;
    struct sockaddr_un sock_addr;
    int                connect_fd;
    struct sockaddr_un connect_addr;
    socklen_t          connect_addr_len = 0;

    listen_fd = socket(PF_UNIX, SOCK_STREAM, 0);
    if (listen_fd < 0) {
        printf("%s: socket create failed", __func__);
        return -1;
    }

    if (verbose)
        printf("%s: listen_fd: %d\n", __func__, listen_fd);

    unlink(socket_name);

    memset(&sock_addr, 0, sizeof(struct sockaddr_un));
    sock_addr.sun_family = AF_UNIX;
    strncpy(sock_addr.sun_path, socket_name, sizeof(sock_addr.sun_path) - 1);

    if (bind(listen_fd, (const struct sockaddr *)&sock_addr, sizeof(struct sockaddr_un))) {
        printf("i%s: bind error", __func__);
        return -1;
    }

    if (listen(listen_fd, 1)) {
        printf("%s: listen error", __func__);
        return -1;
    }

    connect_fd = accept(listen_fd, (struct sockaddr *)&connect_addr, &connect_addr_len);

    if (verbose)
        printf("%s: connect_fd: %d\n", __func__, connect_fd);

    close(listen_fd);
    unlink(socket_name);
    if (connect_fd < 0) {
        printf("%s: accept failed\n", __func__);
        return -1;
    }

    return connect_fd;
}

/* receive a file descriptor from another process.         */
/* Returns the file descriptor.  Note: the integer value   */
/* of the file descriptor may be different from the        */
/* integer value in the other process, but the file        */
/* descriptors in each process will refer to the same file */
/* object in the kernel.                                   */
int EGLStreamReceivefd(int connect_fd)
{
    struct msghdr   msg;
    struct iovec    iov[1];
    char            msg_buf[1];
    char            ctrl_buf[CMSG_SPACE(sizeof(int))];
    struct cmsghdr *cmsg;
    void           *data;
    int             recvfd;

    memset(&msg, 0, sizeof(msg));

    iov[0].iov_base = msg_buf;
    iov[0].iov_len  = sizeof(msg_buf);
    msg.msg_iov     = iov;
    msg.msg_iovlen  = 1;

    msg.msg_control    = ctrl_buf;
    msg.msg_controllen = sizeof(ctrl_buf);

    if (recvmsg(connect_fd, &msg, 0) <= 0) {
        printf("%s: recvmsg failed", __func__);
        return -1;
    }

    cmsg = CMSG_FIRSTHDR(&msg);
    if (!cmsg) {
        printf("%s: NULL message header\n", __func__);
        return -1;
    }
    if (cmsg->cmsg_level != SOL_SOCKET) {
        printf("%s: Message level is not SOL_SOCKET\n", __func__);
        return -1;
    }
    if (cmsg->cmsg_type != SCM_RIGHTS) {
        printf("%s: Message type is not SCM_RIGHTS\n", __func__);
        return -1;
    }

    data   = CMSG_DATA(cmsg);
    recvfd = *(int *)data;

    return recvfd;
}

#endif
