#include <assert.h>
#include <stdio.h>
#include <cuda_runtime.h>

#include "eglapiaccessors.h"
#include "eglframeconsumer.h"

EGLFrameConsumer::EGLFrameConsumer(int fifoLength, int latency)
{
    printf("Initialize EGL frame consumer.\n");

    display = EGL_NO_DISPLAY;
    stream = EGL_NO_STREAM_KHR;

    this->fifoLength = fifoLength;
    if (fifoLength > 0) {
        fifoMode = true;
    } else {
        fifoMode = false;
    }
    this->latency = latency;

    if (!initEGLDisplay()) {
        printf("Cannot initialize EGL display.\n");
        throw;
    }

    if (!initEGLStream()) {
        printf("Cannot initialize EGL Stream.\n");
        throw;
    }

    if (!initEGLCudaConsumer()) {
        printf("Cannot initialize CUDA consumer.\n");
        throw;
    }
}

EGLFrameConsumer::~EGLFrameConsumer()
{
    printf("Destroy EGL frame consumer.\n");

    finalizeEGLCudaConsumer();
    finalizeEGLStream();
}

CUeglFrame* EGLFrameConsumer::fetch()
{
    if (cudaFree(nullptr) != cudaSuccess) {
        printf("Failed to initialize CUDA context");
        return nullptr;
    }

    EGLint streamState = 0;
    if (!eglQueryStreamKHR(display, stream, EGL_STREAM_STATE_KHR, &streamState)) {
        printf("Cuda consumer, eglQueryStreamKHR EGL_STREAM_STATE_KHR failed.\n");
        return nullptr;
    }

    CUresult cuStatus = cuEGLStreamConsumerAcquireFrame(&cudaConnection, &cudaResource, nullptr, 50000);
    if (cuStatus != CUDA_SUCCESS) {
        printf("Cuda Acquire EGL stream frame failed cuStatus=%d.\n", cuStatus);
        return nullptr;
    }

    EGLTimeKHR timeStampNow = 0;
    if (!eglQueryStreamTimeKHR(display, stream, EGL_STREAM_TIME_NOW_KHR, &timeStampNow)) {
        printf("Cuda consumer, eglQueryStreamTimeKHR EGL_STREAM_STATE_KHR failed.\n");
        return nullptr;
    }
    printf("timestamp: %ldms, diff %ldms \n", timeStampNow/1000000, (timeStampNow-timeStamp)/1000000);
    timeStamp = timeStampNow;

    cuStatus = cuGraphicsResourceGetMappedEglFrame(&eglFrame, cudaResource, 0, 0);
    if (cuStatus != CUDA_SUCCESS) {
        printf("Cuda get resource failed with %d.\n", cuStatus);
        cuEGLStreamConsumerReleaseFrame(&cudaConnection, cudaResource, nullptr);
        return nullptr;
    }

    assert(eglFrame.cuFormat == CU_AD_FORMAT_UNSIGNED_INT8);
    assert(eglFrame.eglColorFormat == CU_EGL_COLOR_FORMAT_RGBA);
    assert(eglFrame.planeCount == 1);
    assert(eglFrame.numChannels == 4);
    assert(eglFrame.frameType == CU_EGL_FRAME_TYPE_PITCH);
    assert(eglFrame.pitch == eglFrame.width * 4);
    return &eglFrame;
}

void EGLFrameConsumer::resetStream()
{
    EGLBoolean ret = eglResetStreamNV(display, stream);
    printf("eglResetStreamNV = %d.\n", ret);
}

void EGLFrameConsumer::release()
{
    CUresult cuStatus = cuEGLStreamConsumerReleaseFrame(&cudaConnection, cudaResource, nullptr);
    if (cuStatus != CUDA_SUCCESS) {
        printf("Cuda release frame failed cuStatus=%d.\n", cuStatus);
    }
}

void EGLFrameConsumer::clear()
{
    if (!fifoMode) {
        return;
    }

    if (cudaFree(nullptr) != cudaSuccess) {
        printf("Failed to initialize CUDA context");
        return;
    }

    EGLint streamState = 0;
    if (!eglQueryStreamKHR(display, stream, EGL_STREAM_STATE_KHR, &streamState)) {
        printf("Cuda consumer, eglQueryStreamKHR EGL_STREAM_STATE_KHR failed.\n");
        return;
    }

    if (streamState != EGL_STREAM_STATE_NEW_FRAME_AVAILABLE_KHR) {
        return;
    }

    CUresult cuStatus = cuEGLStreamConsumerAcquireFrame(&cudaConnection, &cudaResource, nullptr, 50000);
    if (cuStatus != CUDA_SUCCESS) {
        printf("Cuda Acquire failed cuStatus=%d.\n", cuStatus);
        return;
    }

    cuStatus = cuEGLStreamConsumerReleaseFrame(&cudaConnection, cudaResource, nullptr);
    if (cuStatus != CUDA_SUCCESS) {
        printf("Cuda release frame failed cuStatus=%d.\n", cuStatus);
    }
}

bool EGLFrameConsumer::initEGLDisplay()
{
    // Obtain the EGL display
    display = EGLDisplayAccessor::getInstance();
    if (display == EGL_NO_DISPLAY) {
        printf("Obtain EGL display failed.\n");
        return false;
    }

    return true;
}

bool EGLFrameConsumer::initEGLStream()
{
    const EGLint streamAttrMailboxMode[] = { /*EGL_SUPPORT_REUSE_NV, EGL_FALSE,*/ EGL_NONE };
    const EGLint streamAttrFIFOMode[] = { EGL_STREAM_FIFO_LENGTH_KHR, fifoLength, EGL_SUPPORT_RESET_NV, EGL_TRUE, EGL_NONE };

    if (!setupEGLExtensions()) {
        return false;
    }

    stream = eglCreateStreamKHR(display, fifoMode ? streamAttrFIFOMode : streamAttrMailboxMode);
    if (stream == EGL_NO_STREAM_KHR) {
        printf("Couldn't create stream.\n");
        return false;
    }

    if (!eglStreamAttribKHR(display, stream, EGL_CONSUMER_LATENCY_USEC_KHR, latency)) {
        printf("Consumer: streamAttribKHR EGL_CONSUMER_LATENCY_USEC_KHR failed.\n");
    }
    if (!eglStreamAttribKHR(display, stream, EGL_CONSUMER_ACQUIRE_TIMEOUT_USEC_KHR, 50000)) {
        printf("Consumer: streamAttribKHR EGL_CONSUMER_ACQUIRE_TIMEOUT_USEC_KHR failed.\n");
    }

    // Get stream attributes
    if (!eglQueryStreamKHR(display, stream, EGL_STREAM_FIFO_LENGTH_KHR, &fifoLength)) {
        printf("Consumer: eglQueryStreamKHR EGL_STREAM_FIFO_LENGTH_KHR failed.\n");
    }
    if (!eglQueryStreamKHR(display, stream, EGL_CONSUMER_LATENCY_USEC_KHR, &latency)) {
        printf("Consumer: eglQueryStreamKHR EGL_CONSUMER_LATENCY_USEC_KHR failed.\n");
    }

    if (fifoMode != (fifoLength > 0)) {
        printf("EGL Stream consumer - Unable to set FIFO mode.\n");
        fifoMode = false;
    }
    if (fifoMode) {
        printf("EGL Stream consumer - Mode: FIFO, Length: %d, latency %d.\n", fifoLength, latency);
    } else {
        printf("EGL Stream consumer - Mode: Mailbox.\n");
    }

    return true;
}

void EGLFrameConsumer::finalizeEGLStream()
{
    if (stream != EGL_NO_STREAM_KHR) {
        eglDestroyStreamKHR(display, stream);
        stream = EGL_NO_STREAM_KHR;
    }

    // acquire and release unprocess frames.
}

bool EGLFrameConsumer::initEGLCudaConsumer()
{
    if (cudaFree(nullptr) != cudaSuccess) {
        printf("Failed to initialize CUDA context.\n");
        return false;
    }

    printf("Connect CUDA consumer.\n");
    CUresult curesult = cuEGLStreamConsumerConnect(&cudaConnection, stream);
    if (curesult != CUDA_SUCCESS) {
        printf("Connect CUDA consumer ERROR %d.\n", curesult);
        return false;
    }

    return true;
}

void EGLFrameConsumer::finalizeEGLCudaConsumer()
{
    if (cudaConnection) {
        if (cudaFree(nullptr) != cudaSuccess) {
            printf("Failed to initialize CUDA context.\n");
            return;
        }

        cuEGLStreamConsumerDisconnect(&cudaConnection);
        cudaConnection = nullptr;
    }
}