#include <fstream>
#include <stdexcept>

#include "OCIOPluginBase.h"
#include "OCIOConfigManager.h"

#include <OpenColorIO/OpenColorIO.h>
#include <OpenColorIO/OpenColorTypes.h>
#include <DDImage/plugins.h>
#include "DDImage/gl.h"
#include "DDImage/ExtendedOpProperties.h"

namespace OCIO = OCIO_NAMESPACE;

using namespace DD::Image;

namespace
{
  /// Cache for storing LUT data which is shared between ops.  The data is ref-counted
  /// to avoid keeping it in memory indefinitely.
  ///
  /// The rationale for this is that on the timeline, each TrackItem has an OCIOColorSpace node.
  /// The chances are that many of the items have the same file type, and thus will use the same color space.
  /// Caches the calls to generateLut1D().
  ///
  /// Access to the cache is synchronised, although I'm not sure if that's necessary.
  class LutImageCache
  {
  public:
    LutImageCache();

    void incRef(const std::string& key);
    void decRef(const std::string& key);
    SamplerTextureStore* findInCache(const std::string& key);

    DD::Image::Lock& lock() { return _lock; }

  private:
    using RefCountDataPair = std::pair< unsigned int, SamplerTextureStore >; //< Data stored in the map, with first being the ref-count, and second the LUT data
    using Map = std::map< std::string, RefCountDataPair >;
    Map _map;
    DD::Image::Lock _lock;
  };

  LutImageCache::LutImageCache()
  {
  }

  /// Add a reference to data in the cache.  If not already present, inserts an entry for the key,
  /// but does not initialise the data.
  void LutImageCache::incRef(const std::string& key)
  {
    DD::Image::Guard lockGuard(_lock);
    // Increment the ref count, if key was not already present in the map, this will create a default-constructed entry
    // with ref-count 0 and an empty vector.  The initialisation of the data is done by OCIOPluginBase::setupGPUShaders.
    _map[key].first += 1;
  }

  /// Remove a reference to data in the cache.  If the ref-count reaches zero, the data is removed from the cache.
  void LutImageCache::decRef(const std::string& key)
  {
    DD::Image::Guard lockGuard(_lock);

    // Search the map
    Map::iterator it = _map.find(key);

    mFnAssertMsg(it != _map.end(), "LutImageCache::decRef key not found");

    if ( it != _map.end() ) {
      it->second.first -= 1; // Decrement the ref-count

      // If the ref-count has reached 0, remove the entry
      if ( it->second.first == 0 ) {
        _map.erase(it);
      }
    }
  }

  /// Fetch the data from the cache
  SamplerTextureStore* LutImageCache::findInCache(const std::string& key)
  {
    DD::Image::Guard lockGuard(_lock);
    Map::iterator it = _map.find(key);
    mFnAssertMsg(it != _map.end(), "Lut1DCache::findInCache key not found");
    if ( it != _map.end() ) {
      return &(it->second.second); // Return the vector data
    }
    else {
      return nullptr;
    }
  }

  static LutImageCache sLutImageCache;
}

const char* OCIOPluginBase::kContextTooltip =
  "\nOCIO Contexts allow you to apply specific LUTs or grades to different shots.\n\n"
  "Here you can specify the context name (key) and its corresponding value.\n\n"
  "Full details of how to set up contexts and add them to your config can be found in the OpenColorIO documentation: \n\n"
  "https://opencolorio.readthedocs.io/en/stable/guides/authoring/contexts.html";

const int OCIOPluginBase::kNumOCIOContextVariables = 4;

OCIOPluginBase::LutCacheReference::LutCacheReference()
{
}

OCIOPluginBase::LutCacheReference::~LutCacheReference()
{
  setKey(std::string()); // Ensure the cache entry is dereferenced.
}

/// Set the cache key and get a reference to it.  If there was a previous key set, decrements its ref-count.
void OCIOPluginBase::LutCacheReference::setKey(const std::string &newKey)
{
  if ( newKey != _key ) {
    // If there was already a key, decrement the refcount
    if ( !_key.empty() ) {
      sLutImageCache.decRef(_key);
    }

    _key = newKey;

    // If new key is not empty, increment the refcount
    if ( !_key.empty() ) {
      sLutImageCache.incRef(_key);
    }
  }
}

/// Get the LUTTextureInfo from the cache.
SamplerTextureStore* OCIOPluginBase::LutCacheReference::lutImageRef()
{
  mFnAssertMsg(!_key.empty(), "LutCacheReference empty key");
  return sLutImageCache.findInCache(_key);
}

const Op::VersionInfo OCIOPluginBase::kOcioVersion { GetOCIOVersion() };

OCIOPluginBase::OCIOPluginBase(Node *node)
: PixelIop(node)
{
  _config = OCIO::ConstConfigRcPtr();
}

OCIOPluginBase::~OCIOPluginBase()
{
  deleteAllTextures();
}

bool OCIOPluginBase::validateConfig()
{
  try {
    OCIO::ConstConfigRcPtr config = getConfigForProject();

    if (config != _config) {
      _config = config;
      configChanged();
      return false;
    }
  }
  catch(OCIO::Exception &e) {
    // Don't need to show the message again ...
    return false;
  }

   return true;
}

// Functions to replace the OCIOv1 getGpuLut3D function, which was used to generate the 3D lut values.
// For a given index, the function will modifed the TextureId references with the values needed to construct
// a LUT sampler texture that will be used by the given OCIO shaderDesc.
//
// outputTexture: the TextureId data.
void OCIOPluginBase::generateLut1D(LUTTextureInfo& outputTexture, unsigned int textureIndex, const OCIO::GpuShaderDescRcPtr shaderDesc)
{
  if (textureIndex < shaderDesc->getNumTextures()) {
    const char *textureName = nullptr;
    const char *samplerName = nullptr;
    unsigned width = 0;
    unsigned height = 0;

    auto textureType = OCIO::GpuShaderDesc::TEXTURE_RGB_CHANNEL;
    OCIO::Interpolation interpolation = OCIO::INTERP_LINEAR;
    shaderDesc->getTexture(textureIndex, textureName, samplerName, width, height, textureType, interpolation);

    if (!textureName || !*textureName || !samplerName || !*samplerName || width == 0) {
      throw OCIO::Exception("Failed to generate 1DLut texture.");;
    }

    const float *values = nullptr;
    shaderDesc->getTextureValues(textureIndex, values);
    if (!values) {
      throw OCIO::Exception("Failed to get 1DLut texture values.");;
    }

    outputTexture.m_uid = textureIndex;
    outputTexture.m_textureName = samplerName;
    outputTexture.m_glTextureType = (height > 1) ? GL_TEXTURE_2D : GL_TEXTURE_1D;
    outputTexture.m_glTextureChannelType = textureType;
    outputTexture.m_textureWidth = width;
    outputTexture.m_textureHeight = height;

    const auto channelCount = textureType == OCIO::GpuShaderDesc::TEXTURE_RGB_CHANNEL ? 3 : 1;
    const auto sampleSize = channelCount * width * height;
    outputTexture.m_pixels.insert(outputTexture.m_pixels.end(), &values[0], &values[sampleSize]);
  }
}

void OCIOPluginBase::generateLut3D(LUTTextureInfo& outputTexture, unsigned int textureIndex, const OCIO::GpuShaderDescRcPtr shaderDesc)
{
  if (textureIndex < shaderDesc->getNum3DTextures()) {
    const char *textureName = nullptr;
    const char *samplerName = nullptr;
    unsigned edgeLen = 0;
    unsigned width = 0;
    unsigned height = 0;

    OCIO::Interpolation interpolation = OCIO::INTERP_LINEAR;
    shaderDesc->get3DTexture(textureIndex, textureName, samplerName, edgeLen, interpolation);

    if (!textureName || !*textureName || !samplerName || !*samplerName || edgeLen == 0) {
      throw OCIO::Exception("Failed to generate 3DLut texture.");;
    }

    const float *values = nullptr;
    shaderDesc->get3DTextureValues(textureIndex, values);
    if (!values) {
      throw OCIO::Exception("Failed to get 3DLut texture values.");;
    }

    outputTexture.m_uid = textureIndex;
    outputTexture.m_textureName = samplerName;
    outputTexture.m_glTextureType = GL_TEXTURE_3D;
    outputTexture.m_glTextureChannelType = OCIO::GpuShaderDesc::TEXTURE_RGB_CHANNEL;
    outputTexture.m_textureWidth = width;
    outputTexture.m_textureHeight = height;
    outputTexture.m_textureEdgeLen = edgeLen;

    const auto channelCount = 3;
    const auto sampleSize = channelCount * edgeLen * edgeLen * edgeLen;
    outputTexture.m_pixels.insert(outputTexture.m_pixels.end(), &values[0], &values[sampleSize]);
  }
}

// Generate the sampler texture data if needed and store them in a list of TextureId
void OCIOPluginBase::generateSamplerTextures(const OCIO::GpuShaderDescRcPtr& shaderDesc)
{
  // With OCIOv2, the color transformation could contain any number of 1D, 2D & 3D LUTs (i.e. GPU textures)
  // which is an important difference with the OCIOv1 GPU code path (where at most one 3D LUT was expected).

  // Check if any 3D textures are needed by the current shaderDesc
  const unsigned maxTexture3D = shaderDesc->getNum3DTextures();

  for (unsigned idx = 0; idx < maxTexture3D; ++idx) {
    std::unique_ptr<LUTTextureInfo> newTexture = std::make_unique<LUTTextureInfo>();

    generateLut3D(*newTexture, idx, shaderDesc);

    m_samplerTextureStore->push_back(std::move(newTexture));
  }

  // Check if any 2D or 1D textures are needed by the current shaderDesc
  const unsigned maxTexture2D = shaderDesc->getNumTextures();

  for (unsigned idx = 0; idx < maxTexture2D; ++idx) {
    std::unique_ptr<LUTTextureInfo> newTexture = std::make_unique<LUTTextureInfo>();

    generateLut1D(*newTexture, idx, shaderDesc);

    m_samplerTextureStore->push_back(std::move(newTexture));
  }
}

void DD::Image::OCIOPluginBase::deleteAllTextures()
{
  if (!m_samplerTextureStore) {
    return;
  }

  for (auto& textureId : *m_samplerTextureStore)
  {
      if (textureId->m_textureHandle != 0) {
        glDeleteTextures(1, &textureId->m_textureHandle);
      }
  }
  m_samplerTextureStore->clear();
}

void OCIOPluginBase::setupGPUShaders(OCIO::ConstProcessorRcPtr &processor)
{
  // gpu path disabled if there is no processor or if the plugin does not support it in this context
  // The Nuke DAG context is disabled because of some weird bug with OCIO where enabling the CPU path
  // the first time still uses software - then enabling and disabling the GPU buffer depth makes it start working
  if (!processor.get() || processor->isNoOp() || !isGPUAccelerated()) {
    m_gpuEngineDecl = "";
    m_gpuEngineBody = "";
    return;
  }

  // Set shader.
  const OCIO::OptimizationFlags optimization{ OCIO::OPTIMIZATION_DEFAULT };

  const auto gpuProcessor = processor->getOptimizedGPUProcessor(optimization);
  const std::string cacheID = gpuProcessor->getCacheID();
  if (cacheID != _gpuProcessorCacheID) {
    _gpuProcessorCacheID = cacheID;

    _shaderDescription = OCIO::GpuShaderDesc::CreateShaderDesc();
    _shaderDescription->setLanguage(OCIO::GPU_LANGUAGE_GLSL_1_2);
    const auto shaderFuncName = std::string( this->Class() ) + "_$$ocio";  // ReduceGPUOps will replace the appended $$ with the appropriate ID.
    _shaderDescription->setFunctionName(shaderFuncName.c_str());
    _shaderDescription->setResourcePrefix("$$nuke");
    gpuProcessor->extractGpuShaderInfo(_shaderDescription);
  }

  // Get the cached LUT textures from the cache if there's one.
  m_lutImageCacheRef.setKey(cacheID);
  m_samplerTextureStore = m_lutImageCacheRef.lutImageRef();

  // If the LUT hasn't yet been initialised, do it now.
  {
    DD::Image::Guard lockGuard( sLutImageCache.lock() );
    if (m_samplerTextureStore && m_samplerTextureStore->empty()) {
      try {
        // For the given _shaderDescription, get the data for any sample textures that will be needed
        // for the GL calls.
        generateSamplerTextures(_shaderDescription);
      }
      catch (OCIO::Exception& e) {
        error(e.what());
        m_gpuEngineDecl = "";
        m_gpuEngineBody = "";
        return;
      }
    }
  }

  // Generate the GLSL code

  // GPU path shader declaration
  std::stringstream gpuDecl;
  gpuDecl << std::endl
          << _shaderDescription->getShaderText()
          << std::endl;

  m_gpuEngineDecl = gpuDecl.str();

  // GPU path shader body
  std::stringstream gpuBody;
  gpuBody << std::endl
          << "OUT.rgb = " << _shaderDescription->getFunctionName() << "(OUT).rgb;" << std::endl
          << std::endl;

  m_gpuEngineBody = gpuBody.str();
}

void OCIOPluginBase::checkGLError(const char* scope)
{
  GLenum glError = glGetError();
  if (glError != GL_NO_ERROR) {
    std::ostringstream msg;
    msg << scope << ": " << gl_errorstring(glError) << std::endl;
    std::cout << msg.str() << std::endl;
    warning(msg.str().c_str());
  }
}

const char* OCIOPluginBase::gpuEngine_decl() const
{
  if (!isGPUAccelerated()) {
    return nullptr;
  }

  return m_gpuEngineDecl.c_str();
}

const char* OCIOPluginBase::gpuEngine_body() const
{
  if (!isGPUAccelerated()) {
    return nullptr;
  }

  return m_gpuEngineBody.c_str();
}

Hash OCIOPluginBase::gpuEngine_shader_hash_at(double time)
{
  return gpuEngine_shader_hash_at_impl(time);
}

int OCIOPluginBase::gpuEngine_getNumRequiredTexUnits() const
{
  if (!isGPUAccelerated()) {
    return 0;
  }

  return m_samplerTextureStore ? static_cast<int>(m_samplerTextureStore->size()) : 1;
}

bool OCIOPluginBase::shouldProcessEngineGL() const
{
  if (!isGPUAccelerated()) {
    return false;
  }

  if (m_gpuEngineDecl.size() == 0)  {
    // nothing to do
    return false;
  }

  if (!m_samplerTextureStore || m_samplerTextureStore->empty()) {
    // No LUT values were created by OCIO shaders generation. So nothing to do.
    return false;
  }

  return true;
}

void OCIOPluginBase::setTextureParams(GLint textureType, GLint textureHandle)
{
  glBindTexture(textureType, textureHandle);
  checkGLError("Binding texture GL begin");

  glPixelStorei(GL_UNPACK_ROW_LENGTH, 0);
  glPixelStorei(GL_UNPACK_ALIGNMENT, 1);

  glTexParameteri(textureType, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
  glTexParameteri(textureType, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
  checkGLError("Setting filter parameters");

  glTexParameteri(textureType, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
  glTexParameteri(textureType, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
  glTexParameteri(textureType, GL_TEXTURE_WRAP_R, GL_CLAMP_TO_EDGE);
  checkGLError("Setting wrap parameters");
}

void DD::Image::OCIOPluginBase::uploadNewTexture(const LUTTextureInfo& textureId)
{
  switch (textureId.m_glTextureType) {

  case GL_TEXTURE_1D:
    glTexImage1D(GL_TEXTURE_1D, 0, textureId.internalTextureGLFormat(),textureId.m_textureWidth,
                 0, textureId.textureGLFormat(), GL_FLOAT, &(textureId.m_pixels[0]));
    checkGLError("Creating 1D texture");
    break;

  case GL_TEXTURE_2D:
    glTexImage2D(GL_TEXTURE_2D, 0, textureId.internalTextureGLFormat(), textureId.m_textureWidth, textureId.m_textureHeight,
                 0, textureId.textureGLFormat(), GL_FLOAT, &(textureId.m_pixels[0]));
    checkGLError("Creating 2D texture");
    break;

  case GL_TEXTURE_3D:
    glTexImage3D(GL_TEXTURE_3D, 0, GL_RGB32F_ARB,
                 textureId.m_textureEdgeLen, textureId.m_textureEdgeLen, textureId.m_textureEdgeLen,
                 0, GL_RGB, GL_FLOAT, &(textureId.m_pixels[0]));
    checkGLError("Creating 3D texture");
    break;

  default:
    error("Invalid GL texture type");

  }
}

bool DD::Image::OCIOPluginBase::createNewTexture(unsigned& textureHandle)
{
  if (textureHandle == 0) {
    glGenTextures(1, &textureHandle);
    checkGLError("Generating texture");
    debug("Processing in GPU mode.");
    return true;
  }

  return false;
}

void OCIOPluginBase::gpuEngine_GL_begin(DD::Image::GPUContext* context)
{
  if (!shouldProcessEngineGL()) {
    return;
  }

  // Cycle through the list of textures needed for the OCIOv2 shader and create them as needed.
  // Then setup and bind the texture for use by the shader program.
  for (auto& textureInfo : *m_samplerTextureStore) {
    // Get the texture unit from Nuke based on it's internal texture tracking
    static_assert(sizeof(GLint) == sizeof(int), "GLint must be same size as int");
    textureInfo->m_textureUnit = context->acquireTextureUnit();

    if (textureInfo->m_textureUnit == -1) {
      throw std::runtime_error("OCIO: Out of GPU texture units.");
    }

    mFnAssertMsg(!textureInfo->m_pixels.empty(), "no lut found");
    if (textureInfo->m_pixels.empty()) {
      throw std::runtime_error("OCIO: No lut found.");
    }

    // Create the GL Texture if it hasn't been created before
    bool newTexture = createNewTexture(textureInfo->m_textureHandle);

    glActiveTextureARB(GL_TEXTURE0_ARB + textureInfo->m_textureUnit);
    checkGLError("Activating texture GL begin");

    // Enable texturing
    setTextureParams(textureInfo->m_glTextureType, textureInfo->m_textureHandle);

    // We'll only need to upload the texture if it didn't exist before
    if (newTexture) {
      uploadNewTexture(*textureInfo);
    }

    // Bind the sampler texture's uniform
    context->bind(textureInfo->m_textureName, textureInfo->m_textureUnit);
    checkGLError("Binding texture via context");

    glPixelStorei(GL_UNPACK_ALIGNMENT, 4); // Set this back to the default
  }
}

void OCIOPluginBase::gpuEngine_GL_end(DD::Image::GPUContext* context)
{
  if (!shouldProcessEngineGL()) {
    return;
  }

  // Walk through the list of textureIds and unbind them
  for (auto& textureInfo : *m_samplerTextureStore) {
    if (textureInfo->m_textureUnit >= 0) {
      glActiveTextureARB(GL_TEXTURE0_ARB + int(textureInfo->m_textureUnit));
      checkGLError("Activating texture GL end");

      glBindTexture(textureInfo->m_glTextureType, 0);
      checkGLError("Unbinding texture GL end");

      context->releaseTextureUnit(textureInfo->m_textureUnit);
      textureInfo->m_textureUnit = -1;
      glActiveTextureARB(GL_TEXTURE0_ARB);
    }
  }
}

void OCIOPluginBase::applyImageCPU(OCIO::ConstCPUProcessorRcPtr &cpuProcessor, OCIO::ImageDesc& imgDesc) const
{
  if (!cpuProcessor) {
    std::string message = std::string(displayName()) + ": OCIO cpu processor not initialised.";
    throw OCIO::Exception(message.c_str());
  }

  cpuProcessor->apply(imgDesc);
}

OCIO::ConstConfigRcPtr OCIOPluginBase::getConfigForProject() const
{
  return NukeOCIOConfigManager::LoadConfig(this);
}

OCIO::ConstContextRcPtr OCIOPluginBase::getLocalContext()
{
  try {
    OCIO::ConstConfigRcPtr config = getConfigForProject();
    OCIO::ConstContextRcPtr context = config->getCurrentContext();
    OCIO::ContextRcPtr mutableContext;

    auto extendedOpProperties = getExtendedOpProperties();
    for (int ctx=0; ctx < extendedOpProperties->numKeyValuePairs(); ctx++)  {
      const auto& currentCtxVariable = extendedOpProperties->keyValuePair(ctx);

      if (!currentCtxVariable.first.empty()) {
        if(!mutableContext) {
          mutableContext = context->createEditableCopy();
        }
        mutableContext->setStringVar(currentCtxVariable.first.c_str(), currentCtxVariable.second.c_str());
      }
    }

    if (mutableContext) {
      context = mutableContext;
    }

    return context;
  }
  catch(OCIO::Exception &e) {
    error(e.what());
    return OCIO::ConstContextRcPtr();
  }
}

void OCIOPluginBase::updateConfigBeforeCreatingKnobs(DD::Image::Knob_Callback f)
{
  if (f.makeKnobs()) {
    configChanged();
  }
}