/*
 * Adapted from OpenColorIO with this license:
 *
 * Copyright (c) 2003-2010 Sony Pictures Imageworks Inc., et al.
 * 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 Sony Pictures Imageworks 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 AND CONTRIBUTORS
 * "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.
 *
 * Modifications Copyright 2013, Blender Foundation.
 */

#include <limits>
#include <sstream>
#include <string.h>

#include "glew-mx.h"

#ifdef _MSC_VER
#  pragma warning(push)
#  pragma warning(disable : 4251 4275)
#endif
#include <OpenColorIO/OpenColorIO.h>
#ifdef _MSC_VER
#  pragma warning(pop)
#endif

extern "C" {
#include "GPU_immediate.h"
}

using namespace OCIO_NAMESPACE;

#include "MEM_guardedalloc.h"

#include "ocio_impl.h"

static const int LUT3D_EDGE_SIZE = 64;
static const int SHADER_CACHE_SIZE = 4;

extern "C" char datatoc_gpu_shader_display_transform_glsl[];
extern "C" char datatoc_gpu_shader_display_transform_vertex_glsl[];

/* **** OpenGL drawing routines using GLSL for color space transform ***** */

typedef struct OCIO_GLSLShader {
  /* Cache ID */
  std::string lut3dCacheID;
  std::string shaderCacheID;

  /* LUT */
  bool lut3d_texture_allocated; /* boolean flag indicating whether
                                 * lut texture is allocated
                                 */
  bool lut3d_texture_valid;

  GLuint lut3d_texture; /* OGL texture ID for 3D LUT */

  float *lut3d; /* 3D LUT table */

  /* Dither */
  bool use_dither;

  /* Curve Mapping */
  bool use_curve_mapping;
  bool curve_mapping_texture_allocated;
  bool curve_mapping_texture_valid;
  GLuint curve_mapping_texture;
  size_t curve_mapping_cache_id;

  /* Alpha Predivide */
  bool use_predivide;

  /* GLSL stuff */
  GLuint ocio_shader;
  GLuint vert_shader;
  GLuint program;
  GPUShaderInterface *shader_interface;
} GLSLDrawState;

typedef struct OCIO_GLSLDrawState {
  /* Shader Cache */
  OCIO_GLSLShader *shader_cache[SHADER_CACHE_SIZE];

  /* Previous OpenGL state. */
  GLint last_texture, last_texture_unit;
} OCIO_GLSLDrawState;

static GLuint compileShaderText(GLenum shaderType, const char *text)
{
  GLuint shader;
  GLint stat;

  shader = glCreateShader(shaderType);
  glShaderSource(shader, 1, (const GLchar **)&text, NULL);
  glCompileShader(shader);
  glGetShaderiv(shader, GL_COMPILE_STATUS, &stat);

  if (!stat) {
    GLchar log[1000];
    GLsizei len;
    glGetShaderInfoLog(shader, 1000, &len, log);
    fprintf(stderr, "Shader compile error:\n%s\n", log);
    return 0;
  }

  return shader;
}

static GLuint linkShaders(GLuint ocio_shader, GLuint vert_shader)
{
  if (!ocio_shader || !vert_shader)
    return 0;

  GLuint program = glCreateProgram();

  glAttachShader(program, ocio_shader);
  glAttachShader(program, vert_shader);

  glLinkProgram(program);

  /* check link */
  {
    GLint stat;
    glGetProgramiv(program, GL_LINK_STATUS, &stat);
    if (!stat) {
      GLchar log[1000];
      GLsizei len;
      glGetProgramInfoLog(program, 1000, &len, log);
      fprintf(stderr, "Shader link error:\n%s\n", log);
      return 0;
    }
  }

  return program;
}

static OCIO_GLSLDrawState *allocateOpenGLState(void)
{
  return (OCIO_GLSLDrawState *)MEM_callocN(sizeof(OCIO_GLSLDrawState), "OCIO OpenGL State struct");
}

/* Ensure LUT texture and array are allocated */
static bool ensureLUT3DAllocated(OCIO_GLSLShader *shader)
{
  int num_3d_entries = 3 * LUT3D_EDGE_SIZE * LUT3D_EDGE_SIZE * LUT3D_EDGE_SIZE;

  if (shader->lut3d_texture_allocated)
    return shader->lut3d_texture_valid;

  glGenTextures(1, &shader->lut3d_texture);

  shader->lut3d = (float *)MEM_callocN(sizeof(float) * num_3d_entries, "OCIO GPU 3D LUT");

  glActiveTexture(GL_TEXTURE1);
  glBindTexture(GL_TEXTURE_3D, shader->lut3d_texture);
  glTexParameteri(GL_TEXTURE_3D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
  glTexParameteri(GL_TEXTURE_3D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
  glTexParameteri(GL_TEXTURE_3D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
  glTexParameteri(GL_TEXTURE_3D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
  glTexParameteri(GL_TEXTURE_3D, GL_TEXTURE_WRAP_R, GL_CLAMP_TO_EDGE);

  /* clean glError buffer */
  while (glGetError() != GL_NO_ERROR) {
  }

  glTexImage3D(GL_TEXTURE_3D,
               0,
               GL_RGB16F_ARB,
               LUT3D_EDGE_SIZE,
               LUT3D_EDGE_SIZE,
               LUT3D_EDGE_SIZE,
               0,
               GL_RGB,
               GL_FLOAT,
               shader->lut3d);

  shader->lut3d_texture_allocated = true;

  /* GL_RGB16F_ARB could be not supported at some drivers
   * in this case we could not use GLSL display
   */
  shader->lut3d_texture_valid = glGetError() == GL_NO_ERROR;

  return shader->lut3d_texture_valid;
}

static bool ensureCurveMappingAllocated(OCIO_GLSLShader *shader,
                                        OCIO_CurveMappingSettings *curve_mapping_settings)
{
  if (shader->curve_mapping_texture_allocated)
    return shader->curve_mapping_texture_valid;

  glGenTextures(1, &shader->curve_mapping_texture);

  glActiveTexture(GL_TEXTURE2);
  glBindTexture(GL_TEXTURE_1D, shader->curve_mapping_texture);
  glTexParameteri(GL_TEXTURE_1D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
  glTexParameteri(GL_TEXTURE_1D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
  glTexParameteri(GL_TEXTURE_1D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
  glTexParameteri(GL_TEXTURE_1D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
  glTexParameteri(GL_TEXTURE_1D, GL_TEXTURE_WRAP_R, GL_CLAMP_TO_EDGE);

  /* clean glError buffer */
  while (glGetError() != GL_NO_ERROR) {
  }

  glTexImage1D(GL_TEXTURE_1D,
               0,
               GL_RGBA16F,
               curve_mapping_settings->lut_size,
               0,
               GL_RGBA,
               GL_FLOAT,
               curve_mapping_settings->lut);

  shader->curve_mapping_texture_allocated = true;

  /* GL_RGB16F_ARB could be not supported at some drivers
   * in this case we could not use GLSL display
   */
  shader->curve_mapping_texture_valid = glGetError() == GL_NO_ERROR;

  return shader->curve_mapping_texture_valid;
}

static void freeGLSLShader(OCIO_GLSLShader *shader)
{
  if (shader->curve_mapping_texture_allocated) {
    glDeleteTextures(1, &shader->curve_mapping_texture);
  }

  if (shader->lut3d_texture_allocated) {
    glDeleteTextures(1, &shader->lut3d_texture);
  }

  if (shader->lut3d) {
    MEM_freeN(shader->lut3d);
  }

  if (shader->program) {
    glDeleteProgram(shader->program);
  }

  if (shader->shader_interface) {
    GPU_shaderinterface_discard(shader->shader_interface);
  }

  if (shader->ocio_shader) {
    glDeleteShader(shader->ocio_shader);
  }

  using std::string;
  shader->lut3dCacheID.~string();
  shader->shaderCacheID.~string();

  MEM_freeN(shader);
}

/* Detect if we can support GLSL drawing */
bool OCIOImpl::supportGLSLDraw()
{
  /* uses GL_RGB16F_ARB */
  return GLEW_VERSION_3_0 || GLEW_ARB_texture_float;
}

/**
 * Setup OpenGL contexts for a transform defined by processor using GLSL
 * All LUT allocating baking and shader compilation happens here.
 *
 * Once this function is called, callee could start drawing images
 * using regular 2D texture.
 *
 * When all drawing is finished, finishGLSLDraw shall be called to
 * restore OpenGL context to it's pre-GLSL draw state.
 */
bool OCIOImpl::setupGLSLDraw(OCIO_GLSLDrawState **state_r,
                             OCIO_ConstProcessorRcPtr *processor,
                             OCIO_CurveMappingSettings *curve_mapping_settings,
                             float dither,
                             bool use_predivide)
{
  ConstProcessorRcPtr ocio_processor = *(ConstProcessorRcPtr *)processor;
  bool use_curve_mapping = curve_mapping_settings != NULL;
  bool use_dither = dither > std::numeric_limits<float>::epsilon();

  /* Create state if needed. */
  OCIO_GLSLDrawState *state;
  if (!*state_r)
    *state_r = allocateOpenGLState();
  state = *state_r;

  glGetIntegerv(GL_TEXTURE_BINDING_2D, &state->last_texture);
  glGetIntegerv(GL_ACTIVE_TEXTURE, &state->last_texture_unit);

  /* Compute cache IDs. */
  GpuShaderDesc shaderDesc;
  shaderDesc.setLanguage(GPU_LANGUAGE_GLSL_1_3);
  shaderDesc.setFunctionName("OCIODisplay");
  shaderDesc.setLut3DEdgeLen(LUT3D_EDGE_SIZE);

  std::string lut3dCacheID = ocio_processor->getGpuLut3DCacheID(shaderDesc);
  std::string shaderCacheID = ocio_processor->getGpuShaderTextCacheID(shaderDesc);

  /* Find matching cached shader. */
  OCIO_GLSLShader *shader = NULL;
  for (int i = 0; i < SHADER_CACHE_SIZE; i++) {
    OCIO_GLSLShader *cached_shader = state->shader_cache[i];
    if (cached_shader == NULL) {
      continue;
    }

    if (cached_shader->lut3dCacheID == lut3dCacheID &&
        cached_shader->shaderCacheID == shaderCacheID &&
        cached_shader->use_predivide == use_predivide &&
        cached_shader->use_curve_mapping == use_curve_mapping &&
        cached_shader->use_dither == use_dither) {
      /* LRU cache, so move to front. */
      for (int j = i; j > 0; j--) {
        state->shader_cache[j] = state->shader_cache[j - 1];
      }
      state->shader_cache[0] = cached_shader;

      shader = cached_shader;
      break;
    }
  }

  if (shader == NULL) {
    /* LRU cache, shift other items back so we can insert at the front. */
    OCIO_GLSLShader *last_shader = state->shader_cache[SHADER_CACHE_SIZE - 1];
    if (last_shader) {
      freeGLSLShader(last_shader);
    }
    for (int j = SHADER_CACHE_SIZE - 1; j > 0; j--) {
      state->shader_cache[j] = state->shader_cache[j - 1];
    }

    /* Allocate memory for shader. */
    shader = (OCIO_GLSLShader *)MEM_callocN(sizeof(OCIO_GLSLShader), "OCIO GLSL Shader");
    state->shader_cache[0] = shader;

    new (&shader->lut3dCacheID) std::string();
    new (&shader->shaderCacheID) std::string();

    shader->lut3dCacheID = lut3dCacheID;
    shader->shaderCacheID = shaderCacheID;
    shader->use_curve_mapping = use_curve_mapping;
    shader->use_dither = use_dither;
    shader->use_predivide = use_predivide;

    bool valid = true;

    /* Compute 3D LUT. */
    if (valid && ensureLUT3DAllocated(shader)) {
      ocio_processor->getGpuLut3D(shader->lut3d, shaderDesc);

      glActiveTexture(GL_TEXTURE1);
      glBindTexture(GL_TEXTURE_3D, shader->lut3d_texture);
      glTexSubImage3D(GL_TEXTURE_3D,
                      0,
                      0,
                      0,
                      0,
                      LUT3D_EDGE_SIZE,
                      LUT3D_EDGE_SIZE,
                      LUT3D_EDGE_SIZE,
                      GL_RGB,
                      GL_FLOAT,
                      shader->lut3d);
    }
    else {
      valid = false;
    }

    /* Allocate curve mapping texture. */
    if (valid && use_curve_mapping) {
      if (!ensureCurveMappingAllocated(shader, curve_mapping_settings)) {
        valid = false;
      }
    }

    if (valid) {
      /* Vertex shader */
      std::ostringstream osv;

      osv << "#version 330\n";
      osv << datatoc_gpu_shader_display_transform_vertex_glsl;

      shader->vert_shader = compileShaderText(GL_VERTEX_SHADER, osv.str().c_str());

      /* Fragment shader */
      std::ostringstream os;

      os << "#version 330\n";

      /* Work around OpenColorIO not supporting latest GLSL yet. */
      os << "#define texture2D texture\n";
      os << "#define texture3D texture\n";

      if (use_predivide) {
        os << "#define USE_PREDIVIDE\n";
      }

      if (use_dither) {
        os << "#define USE_DITHER\n";
      }

      if (use_curve_mapping) {
        os << "#define USE_CURVE_MAPPING\n";
      }

      os << ocio_processor->getGpuShaderText(shaderDesc) << "\n";
      os << datatoc_gpu_shader_display_transform_glsl;

      shader->ocio_shader = compileShaderText(GL_FRAGMENT_SHADER, os.str().c_str());

      /* Program */
      if (shader->ocio_shader && shader->vert_shader) {
        shader->program = linkShaders(shader->ocio_shader, shader->vert_shader);
      }

      if (shader->program) {
        if (shader->shader_interface) {
          GPU_shaderinterface_discard(shader->shader_interface);
        }
        shader->shader_interface = GPU_shaderinterface_create(shader->program);
      }
    }
  }

  /* Update curve mapping texture. */
  if (use_curve_mapping && shader->curve_mapping_texture_allocated) {
    if (shader->curve_mapping_cache_id != curve_mapping_settings->cache_id) {
      glActiveTexture(GL_TEXTURE2);
      glBindTexture(GL_TEXTURE_1D, shader->curve_mapping_texture);
      glTexSubImage1D(GL_TEXTURE_1D,
                      0,
                      0,
                      curve_mapping_settings->lut_size,
                      GL_RGBA,
                      GL_FLOAT,
                      curve_mapping_settings->lut);
    }
  }

  /* Bind Shader. */
  if (shader->program) {
    glActiveTexture(GL_TEXTURE1);
    glBindTexture(GL_TEXTURE_3D, shader->lut3d_texture);

    if (use_curve_mapping) {
      glActiveTexture(GL_TEXTURE2);
      glBindTexture(GL_TEXTURE_1D, shader->curve_mapping_texture);
    }

    glActiveTexture(GL_TEXTURE0);

    /* IMM needs vertex format even if we don't draw with it.
     *
     * NOTE: The only reason why it's here is because of Cycles viewport.
     * All other areas are managing their own vertex formats.
     * Doing it here is probably harmless, but kind of stupid.
     *
     * TODO(sergey): Look into some nicer solution.
     */
    GPUVertFormat *format = immVertexFormat();
    GPU_vertformat_attr_add(format, "pos", GPU_COMP_F32, 2, GPU_FETCH_FLOAT);
    GPU_vertformat_attr_add(format, "texCoord", GPU_COMP_F32, 2, GPU_FETCH_FLOAT);
    immBindProgram(shader->program, shader->shader_interface);

    immUniform1i("image_texture", 0);
    immUniform1i("lut3d_texture", 1);

    if (use_dither) {
      immUniform1f("dither", dither);
    }

    if (use_curve_mapping) {
      immUniform1i("curve_mapping_texture", 2);
      immUniform1i("curve_mapping_lut_size", curve_mapping_settings->lut_size);
      immUniform4iv("use_curve_mapping_extend_extrapolate",
                    curve_mapping_settings->use_extend_extrapolate);
      immUniform4fv("curve_mapping_mintable", curve_mapping_settings->mintable);
      immUniform4fv("curve_mapping_range", curve_mapping_settings->range);
      immUniform4fv("curve_mapping_ext_in_x", curve_mapping_settings->ext_in_x);
      immUniform4fv("curve_mapping_ext_in_y", curve_mapping_settings->ext_in_y);
      immUniform4fv("curve_mapping_ext_out_x", curve_mapping_settings->ext_out_x);
      immUniform4fv("curve_mapping_ext_out_y", curve_mapping_settings->ext_out_y);
      immUniform4fv("curve_mapping_first_x", curve_mapping_settings->first_x);
      immUniform4fv("curve_mapping_first_y", curve_mapping_settings->first_y);
      immUniform4fv("curve_mapping_last_x", curve_mapping_settings->last_x);
      immUniform4fv("curve_mapping_last_y", curve_mapping_settings->last_y);
      immUniform3fv("curve_mapping_black", curve_mapping_settings->black);
      immUniform3fv("curve_mapping_bwmul", curve_mapping_settings->bwmul);
    }

    return true;
  }
  else {
    glActiveTexture(state->last_texture_unit);
    glBindTexture(GL_TEXTURE_2D, state->last_texture);

    return false;
  }
}

void OCIOImpl::finishGLSLDraw(OCIO_GLSLDrawState *state)
{
  glActiveTexture(state->last_texture_unit);
  glBindTexture(GL_TEXTURE_2D, state->last_texture);
  immUnbindProgram();
}

void OCIOImpl::freeGLState(OCIO_GLSLDrawState *state)
{
  for (int i = 0; i < SHADER_CACHE_SIZE; i++) {
    if (state->shader_cache[i]) {
      freeGLSLShader(state->shader_cache[i]);
    }
  }

  MEM_freeN(state);
}
