/*
 * Copyright (c) 2013, NVIDIA CORPORATION. All rights reserved.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a
 * copy of this software and associated documentation files (the "Software"),
 * to deal in the Software without restriction, including without limitation
 * the rights to use, copy, modify, merge, publish, distribute, sublicense,
 * and/or sell copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
 * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 * DEALINGS IN THE SOFTWARE.
 */

#include "shared.h"
#include "create_window.h"

#define NUM_FRAMES           100
#define WINDOW_WIDTH         600
#define WINDOW_HEIGHT        600
#define STREAM_WIDTH         128
#define STREAM_HEIGHT        128

// Shared.
static EGLDisplay dpy;
static EGLStreamKHR stream;
static int frame = 0;

// Consumer.
static EGLNativeWindowType consumerWin;
static EGLContext consumerCtx;
static EGLSurface consumerSurf;
static GLuint consumerRotationMatrixUniform;

// Producer.
static EGLContext producerCtx;
static EGLSurface producerSurf;
static GLuint producerScaleUniform;

static void DoTest(void)
{
    while (frame++ < NUM_FRAMES) {

        // Producer rendering.
        {
            float scale = (float)((frame % 50) + 50) / 100.0f;

            if (!eglMakeCurrent(dpy, producerSurf, producerSurf, producerCtx))
                FATAL_ERROR("Could not make producer context current");

            glClear(GL_COLOR_BUFFER_BIT);
            glUniform1f(producerScaleUniform, scale);
            glDrawArrays(GL_TRIANGLE_STRIP, 0, 3);
            glFinish();

            if (!eglSwapBuffers(dpy, producerSurf))
                FATAL_ERROR("SwapBuffers failed");
        }

        // Consumer rendering.
        {
            static const GLushort cubeIndices[] = {
                0, 1, 2, 3, 7, 1, 5, 4, 7, 6, 2, 4, 0, 1
            };

            if (!eglMakeCurrent(dpy, consumerSurf, consumerSurf, consumerCtx))
                FATAL_ERROR("Could not make consumer context current");

            if (!eglStreamConsumerAcquireKHR(dpy, stream))
                FATAL_ERROR("Failed to acquire frame");

            GLfloat rotMat[16];
            CreateRotationMatrix(rotMat, (float)frame, 0.0f, 1.0f, 0.0f);
            glUniformMatrix4fv(consumerRotationMatrixUniform, 1, GL_FALSE, rotMat);
            glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
            glDrawElements(GL_TRIANGLE_STRIP, 14, GL_UNSIGNED_SHORT, cubeIndices);
            glFinish();

            if (!eglSwapBuffers(dpy, consumerSurf))
                FATAL_ERROR("Failed in eglSwapBuffers");
        }
    }
}

static void InitializeEgl(void)
{
    InitEGLExtFuncs();

    dpy = eglGetDisplay(EGL_DEFAULT_DISPLAY);
    if (dpy == EGL_NO_DISPLAY)
        FATAL_ERROR("Could not get EGL display");
    if (!eglInitialize(dpy, 0, 0))
        FATAL_ERROR("Could not initialize EGL display");
}

static void InitializeStream(void)
{
    static EGLint streamAttrs[] = {
        EGL_NONE, EGL_NONE,
        EGL_NONE
    };
    stream = eglCreateStreamKHR(dpy, streamAttrs);
    if (stream == EGL_NO_STREAM_KHR)
        FATAL_ERROR("Error creating stream");
}

static void InitializeConsumer(void)
{
    static const char vertexShaderSource[] =
        "attribute vec4 aPosition;\n"
        "attribute vec2 aTexCoord;\n"
        "varying vec2 vTexCoord;\n"
        "uniform mat4 uProjectionMat;\n"
        "uniform mat4 uRotationMat;\n"
        "void main() {\n"
        "    vec4 pos = uRotationMat * aPosition;\n"
        "    pos.z -= 5.0;\n"
        "    gl_Position = uProjectionMat * pos;\n"
        "    vTexCoord = aTexCoord;\n"
        "}\n";

    static const char fragmentShaderSource[] =
        "#extension GL_OES_EGL_image_external : enable\n"
        "precision lowp float;\n"
        "uniform samplerExternalOES uTex;\n"
        "varying vec2 vTexCoord;\n"
        "void main() {\n"
        "    gl_FragColor = texture2D(uTex, vec2(vTexCoord.x, 1.0-vTexCoord.y));\n"
        "    if (vTexCoord.x < 0.01 || vTexCoord.x > 0.99 || vTexCoord.y < 0.01 || vTexCoord.y > 0.99) {\n"
        "        gl_FragColor = vec4(0.0, 0.5, 0.0, 1.0);\n"
        "    }\n"
        "}\n";

    static const GLfloat cubeVertices[] = {
        -1.0, -1.0,  1.0, 1.0, -1.0,  1.0, -1.0,  1.0,  1.0, 1.0,  1.0,  1.0,
        -1.0, -1.0, -1.0, 1.0, -1.0, -1.0, -1.0,  1.0, -1.0, 1.0,  1.0, -1.0,
    };

    static const GLfloat cubeTexCoords[] = {
        1.0, 0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0,
        0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0,
        1.0, 1.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0,
        0.0, 1.0, 1.0, 1.0,
    };

    static EGLint configAttrs[] = {
        EGL_RENDERABLE_TYPE, EGL_OPENGL_ES2_BIT,
        EGL_DEPTH_SIZE, 16,
        EGL_NONE
    };

    static EGLint contextAttrs[] = {
        EGL_CONTEXT_CLIENT_VERSION, 2,
        EGL_NONE
    };

    // Create window.
    consumerWin = CreateWindow(WINDOW_WIDTH, WINDOW_HEIGHT, "EGL Streams");

    // Find config.
    EGLint ncfg = 0;
    EGLConfig cfg;
    if (!eglChooseConfig(dpy, configAttrs, &cfg, 1, &ncfg))
        FATAL_ERROR("Could not find EGL config");
    if (ncfg == 0)
        FATAL_ERROR("No EGL config found");

    // Create context.
    consumerCtx = eglCreateContext(dpy, cfg, EGL_NO_CONTEXT, contextAttrs);
    if (consumerCtx == EGL_NO_CONTEXT)
        FATAL_ERROR("Could not create EGL context");

    // Create window surface.
    consumerSurf = eglCreateWindowSurface(dpy, cfg, consumerWin, 0);
    if (consumerSurf == EGL_NO_SURFACE)
        FATAL_ERROR("Could not create EGLSurface for window");

    // Make current.
    if (!eglMakeCurrent(dpy, consumerSurf, consumerSurf, consumerCtx))
        FATAL_ERROR("Could not make context current");

    // Set clear color/viewport.
    glClearColor(0.5f, 0.5f, 0.5f, 1.0f);
    glViewport(10, 10, WINDOW_WIDTH - 20, WINDOW_HEIGHT - 20);
    glEnable(GL_DEPTH_TEST);

    // Set vertex state and create program.
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 0, cubeVertices);
    glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE, 0, cubeTexCoords);
    glEnableVertexAttribArray(0);
    glEnableVertexAttribArray(1);
    GLuint vs = glCreateShader(GL_VERTEX_SHADER);
    GLuint fs = glCreateShader(GL_FRAGMENT_SHADER);
    LoadShader(vs, vertexShaderSource, sizeof(vertexShaderSource));
    LoadShader(fs, fragmentShaderSource, sizeof(fragmentShaderSource));
    GLuint program = glCreateProgram();
    glAttachShader(program, vs);
    glAttachShader(program, fs);
    glBindAttribLocation(program, 0, "aPosition");
    glBindAttribLocation(program, 1, "aTexCoord");
    glLinkProgram(program);
    glUseProgram(program);

    // Set the uniforms.
    GLfloat projMat[16];
    CreatePerspectiveMatrix(projMat, 45.0f, (float)WINDOW_WIDTH/(float)WINDOW_HEIGHT, 1.0f, 1000.0f);
    glUniformMatrix4fv(glGetUniformLocation(program, "uProjectionMat"), 1, GL_FALSE, projMat);
    glUniform1i(glGetUniformLocation(program, "uTex"), 0);
    consumerRotationMatrixUniform = glGetUniformLocation(program, "uRotationMat");

    // Connect a new external texture as a consumer to the stream.
    GLuint tex;
    glGenTextures(1, &tex);
    glBindTexture(GL_TEXTURE_EXTERNAL_OES, tex);
    if (!eglStreamConsumerGLTextureExternalKHR(dpy, stream))
        FATAL_ERROR("Failed to connect consumer");

    if (glGetError() != GL_NO_ERROR)
        FATAL_ERROR("GL Error while initializing consumer");
}

static void InitializeProducer(void)
{
    static const char vertexShaderSource[] =
        "attribute vec2 aPosition;\n"
        "uniform float uScale;\n"
        "void main() {\n"
        "    gl_Position = vec4(aPosition * uScale, 0.0, 1.0);\n"
        "}\n";

    static const char fragmentShaderSource[] =
        "precision lowp float;\n"
        "void main() {\n"
        "    gl_FragColor = vec4(0.0, 1.0, 0.0, 1.0);\n"
        "}\n";

    static const GLfloat triVertices[] = {
        -0.5, -0.5,
         0.5, -0.5,
         0.0,  0.5};

    static EGLint configAttrs[] = {
        EGL_SURFACE_TYPE, EGL_STREAM_BIT_KHR,
        EGL_RENDERABLE_TYPE, EGL_OPENGL_ES2_BIT,
        EGL_NONE
    };

    static EGLint contextAttrs[] = {
        EGL_CONTEXT_CLIENT_VERSION, 2,
        EGL_NONE
    };

    static EGLint surfaceAttrs[] = {
        EGL_WIDTH, STREAM_WIDTH,
        EGL_HEIGHT, STREAM_HEIGHT,
        EGL_NONE
    };

    // Find producer config.
    EGLint ncfg = 0;
    EGLConfig cfg;
    if (!eglChooseConfig(dpy, configAttrs, &cfg, 1, &ncfg))
        FATAL_ERROR("Could not find EGL config for producer");
    if (ncfg == 0)
        FATAL_ERROR("No EGL config found for producer");

    // Create producer context.
    producerCtx = eglCreateContext(dpy, cfg, EGL_NO_CONTEXT, contextAttrs);
    if (producerCtx == EGL_NO_CONTEXT)
        FATAL_ERROR("Could not create producer EGL context");

    // Create producer surface.
    producerSurf = eglCreateStreamProducerSurfaceKHR(dpy, cfg, stream, surfaceAttrs);
    if (producerSurf == EGL_NO_STREAM_KHR)
        FATAL_ERROR("Could not create producer surface");

    // Make producer context current.
    if (!eglMakeCurrent(dpy, producerSurf, producerSurf, producerCtx))
        FATAL_ERROR("Could not make producer context current");

    // Set the clear color/viewport.
    glClearColor(1.0f, 0.0f, 0.0f, 1.0f);
    glViewport(0, 0, STREAM_WIDTH, STREAM_HEIGHT);

    // Set vertex state and create program.
    glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, 0, triVertices);
    glEnableVertexAttribArray(0);
    GLuint vs = glCreateShader(GL_VERTEX_SHADER);
    GLuint fs = glCreateShader(GL_FRAGMENT_SHADER);
    LoadShader(vs, vertexShaderSource, sizeof(vertexShaderSource));
    LoadShader(fs, fragmentShaderSource, sizeof(fragmentShaderSource));
    GLuint program = glCreateProgram();
    glAttachShader(program, vs);
    glAttachShader(program, fs);
    glBindAttribLocation(program, 0, "aPosition");
    glLinkProgram(program);
    glUseProgram(program);
    producerScaleUniform = glGetUniformLocation(program, "uScale");

    if (glGetError() != GL_NO_ERROR)
        FATAL_ERROR("GL Error while initializing consumer");
}

static void TerminateEgl(void)
{
    eglTerminate(dpy);
}

static void TerminateStream(void)
{
    eglDestroyStreamKHR(dpy, stream);
}

static void TerminateConsumer(void)
{
    eglDestroyContext(dpy, consumerCtx);
    eglDestroySurface(dpy, consumerSurf);
    DestroyWindow(consumerWin);
}

static void TerminateProducer(void)
{
    eglDestroyContext(dpy, producerCtx);
    eglDestroySurface(dpy, producerSurf);
}

int main(int argc, char *argv[])
{
    InitializeEgl();
    InitializeStream();
    InitializeConsumer();
    InitializeProducer();
    DoTest();
    TerminateProducer();
    TerminateConsumer();
    TerminateStream();
    TerminateEgl();
    return 0;
}