/* SPDX-FileCopyrightText: 2025 Blender Authors
 *
 * SPDX-License-Identifier: GPL-2.0-or-later */

#include "OCIO_gpu_shader_binder.hh"

#include <sstream>

#include "MEM_guardedalloc.h"

#include "BLI_assert.h"
#include "BLI_build_config.h"
#include "BLI_string_utils.hh"
#include "BLI_vector.hh"

#include "BKE_colortools.hh"

#include "GPU_immediate.hh"
#include "GPU_shader.hh"
#include "GPU_texture.hh"
#include "GPU_uniform_buffer.hh"

#include "gpu_shader_create_info.hh"

#include "OCIO_config.hh"

#include "gpu_shader_binder_internal.hh"
#include "source_processor.hh"
#include "white_point.hh"

namespace blender::ocio {
namespace internal {
/* -------------------------------------------------------------------- */
/** \name GPUTextures
 * \{ */

GPUTextures::~GPUTextures()
{
  for (GPULutTexture &lut : luts) {
    GPU_texture_free(lut.texture);
  }
  if (dummy) {
    GPU_texture_free(dummy);
  }
  if (uniforms_buffer) {
    GPU_uniformbuf_free(uniforms_buffer);
  }
}

bool GPUTextures::initialize_common()
{
  dummy = GPU_texture_create_error(2, false);

  luts.clear();
  uniforms.clear();

  return (dummy != nullptr);
}

/** \} */

/* -------------------------------------------------------------------- */
/** \name GPUTextures
 * \{ */

GPUCurveMappping::~GPUCurveMappping()
{
  if (lut) {
    MEM_freeN(lut);
  }

  if (texture) {
    GPU_texture_free(texture);
  }
  if (buffer) {
    GPU_uniformbuf_free(buffer);
  }
}

void GPUCurveMappping::rasterize(CurveMapping &curve_mapping)
{
  if (lut) {
    MEM_freeN(lut);
  }

  BKE_curvemapping_init(&curve_mapping);
  BKE_curvemapping_premultiply(&curve_mapping, false);
  BKE_curvemapping_table_RGBA(&curve_mapping, &lut, &lut_size);
}

bool GPUCurveMappping::initialize_common(const bool use_curve_mapping)
{
  if (!use_curve_mapping) {
    return true;
  }

  texture = GPU_texture_create_1d(
      "OCIOCurveMap", lut_size, 1, GPU_RGBA16F, GPU_TEXTURE_USAGE_SHADER_READ, nullptr);
  GPU_texture_filter_mode(texture, false);
  GPU_texture_extend_mode(texture, GPU_SAMPLER_EXTEND_MODE_EXTEND);

  buffer = GPU_uniformbuf_create(sizeof(OCIO_GPUCurveMappingParameters));

  if (texture == nullptr || buffer == nullptr) {
    return false;
  }

  return true;
}

/** \} */

/* -------------------------------------------------------------------- */
/** \name GPUDisplayShader
 * \{ */

GPUDisplayShader::~GPUDisplayShader()
{
  if (shader) {
    GPU_shader_free(shader);
  }
  if (parameters_buffer) {
    GPU_uniformbuf_free(parameters_buffer);
  }
}

bool GPUDisplayShader::matches(const GPUDisplayParameters &display_parameters) const
{
  const bool use_curve_mapping = (display_parameters.curve_mapping != nullptr);
  return (this->from_colorspace == display_parameters.from_colorspace &&
          this->view == display_parameters.view && this->display == display_parameters.display &&
          this->look == display_parameters.look && this->use_curve_mapping == use_curve_mapping);
}

bool GPUDisplayShader::initialize_common()
{
  if (!textures.initialize_common()) {
    is_valid = false;
    return false;
  }

  if (!curve_mapping.initialize_common(use_curve_mapping)) {
    is_valid = false;
    return false;
  }

  return true;
}

/** \} */

/* -------------------------------------------------------------------- */
/** \name GPUShaderCache
 * \{ */

GPUShaderCache::~GPUShaderCache()
{
  clear();
}

GPUDisplayShader *GPUShaderCache::get(const GPUDisplayParameters &display_parameters)
{
  for (std::list<GPUDisplayShader>::iterator it = cache_.begin(); it != cache_.end(); it++) {
    if (it->matches(display_parameters)) {
      /* Move to front of the cache to mark as most recently used. */
      if (it != cache_.begin()) {
        cache_.splice(cache_.begin(), cache_, it);
      }
      return &(*it);
    }
  }
  return nullptr;
}

GPUDisplayShader &GPUShaderCache::create_default()
{
  /* Remove least recently used element from cache. */
  while (cache_.size() >= MAX_SIZE) {
    cache_.pop_back();
  }

  /* Create GPU shader. */
  cache_.emplace_front();

  return cache_.front();
}

void GPUShaderCache::clear()
{
  cache_.clear();
}

/** \} */

}  // namespace internal

namespace {

/* -------------------------------------------------------------------- */
/** \name Internal implementation
 * \{ */

/**
 * Process the generated source code, doing necessary tweaks to get it compiled on the current
 * backend.
 * This function will do some code-level adjustments on the code generated by the OpenColorIO and
 * solve all known compatibility issues.
 */
static void process_source(std::string &source)
{
  source_comment_out_uniforms(source);
  source = GPU_shader_preprocess_source(source);

  /* Comparison operator in Metal returns per-element comparison and returns a vector of booleans.
   * Need a special syntax to see if two vec3 are matched.
   *
   * NOTE: The replacement is optimized for transforming code  generated by
   * GradingPrimaryTransform. A more general approach is possible, but for now prefer processing
   * speed.
   *
   * NOTE: The syntax works for all backends Blender supports. */
  BLI_string_replace(
      source, "if ( gamma != vec3(1., 1., 1.) )", "if (! all(equal(gamma, vec3(1., 1., 1.))) )");
}

static void gpu_curve_mapping_update(internal::GPUCurveMappping &gpu_curve_mapping,
                                     CurveMapping &curve_mapping)
{
  /* Test if we need to update. */
  /* TODO(sergey): Use more reliable cache identifier.
   * Something like monotonously incrementing change counter feels to have less collisions. */
  const size_t cache_id = size_t(&curve_mapping) + curve_mapping.changed_timestamp;
  if (gpu_curve_mapping.cache_id == cache_id) {
    return;
  }
  gpu_curve_mapping.cache_id = cache_id;

  /* Update texture. */
  const int offset[3] = {0, 0, 0};
  const int extent[3] = {gpu_curve_mapping.lut_size, 0, 0};
  GPU_texture_update_sub(gpu_curve_mapping.texture,
                         GPU_DATA_FLOAT,
                         gpu_curve_mapping.lut,
                         UNPACK3(offset),
                         UNPACK3(extent));

  /* Update uniforms. */
  OCIO_GPUCurveMappingParameters data;
  for (int i = 0; i < 4; i++) {
    const CurveMap &curve_map = curve_mapping.cm[i];
    data.range[i] = curve_map.range;
    data.mintable[i] = curve_map.mintable;
    data.ext_in_x[i] = curve_map.ext_in[0];
    data.ext_in_y[i] = curve_map.ext_in[1];
    data.ext_out_x[i] = curve_map.ext_out[0];
    data.ext_out_y[i] = curve_map.ext_out[1];
    data.first_x[i] = curve_map.table[0].x;
    data.first_y[i] = curve_map.table[0].y;
    data.last_x[i] = curve_map.table[CM_TABLE].x;
    data.last_y[i] = curve_map.table[CM_TABLE].y;
  }
  for (int i = 0; i < 3; i++) {
    data.black[i] = curve_mapping.black[i];
    data.bwmul[i] = curve_mapping.bwmul[i];
  }
  data.lut_size = gpu_curve_mapping.lut_size;
  data.use_extend_extrapolate = (curve_mapping.flag & CUMA_EXTEND_EXTRAPOLATE) != 0;

  GPU_uniformbuf_update(gpu_curve_mapping.buffer, &data);
}

static void gpu_display_shader_parameters_update(internal::GPUDisplayShader &display_shader,
                                                 const GPUDisplayParameters &display_parameters,
                                                 float4x4 scene_linear_matrix)
{
  bool do_update = false;
  if (display_shader.parameters_buffer == nullptr) {
    display_shader.parameters_buffer = GPU_uniformbuf_create(sizeof(OCIO_GPUParameters));
    do_update = true;
  }

  OCIO_GPUParameters &data = display_shader.parameters;
  if (data.scene_linear_matrix != scene_linear_matrix) {
    data.scene_linear_matrix = scene_linear_matrix;
    do_update = true;
  }
  if (data.exponent != display_parameters.exponent) {
    data.exponent = display_parameters.exponent;
    do_update = true;
  }
  if (data.dither != display_parameters.dither) {
    data.dither = display_parameters.dither;
    do_update = true;
  }
  if (bool(data.use_predivide) != display_parameters.use_predivide) {
    data.use_predivide = display_parameters.use_predivide;
    do_update = true;
  }
  if (bool(data.do_overlay_merge) != display_parameters.do_overlay_merge) {
    data.do_overlay_merge = display_parameters.do_overlay_merge;
    do_update = true;
  }
  if (bool(data.use_hdr) != display_parameters.use_hdr) {
    data.use_hdr = display_parameters.use_hdr;
    do_update = true;
  }

  if (do_update) {
    GPU_uniformbuf_update(display_shader.parameters_buffer, &data);
  }
}

/* Bind the shader and update parameters and uniforms. */
static bool gpu_shader_bind(const Config &config,
                            internal::GPUDisplayShader &display_shader,
                            const GPUDisplayParameters &display_parameters)
{
  using internal::TextureSlot;
  using internal::UniformBufferSlot;

  /* Verify the shader is valid. */
  if (!display_shader.is_valid) {
    return false;
  }

  /* Update and bind curve mapping data. */
  if (display_parameters.curve_mapping) {
    gpu_curve_mapping_update(display_shader.curve_mapping, *display_parameters.curve_mapping);
    GPU_uniformbuf_bind(display_shader.curve_mapping.buffer, UniformBufferSlot::CURVE_MAPPING);
    GPU_texture_bind(display_shader.curve_mapping.texture, TextureSlot::CURVE_MAPPING);
    /* TODO(sergey): Can free the curve mapping's lookup table.
     * Seems minor, maybe not that important. */
  }

  /* Bind textures to sampler units. Texture 0 is set by caller.
   * Uniforms have already been set for texture bind points. */
  if (!display_parameters.do_overlay_merge) {
    /* Avoid missing binds. */
    GPU_texture_bind(display_shader.textures.dummy, TextureSlot::OVERLAY);
  }
  for (int i = 0; i < display_shader.textures.luts.size(); i++) {
    GPU_texture_bind(display_shader.textures.luts[i].texture, TextureSlot::LUTS_OFFSET + i);
  }

  if (display_shader.textures.uniforms_buffer) {
    GPU_uniformbuf_bind(display_shader.textures.uniforms_buffer, UniformBufferSlot::LUTS);
  }

  float3x3 matrix = float3x3::identity() * display_parameters.scale;
  if (display_parameters.use_white_balance) {
    matrix *= calculate_white_point_matrix(
        config, display_parameters.temperature, display_parameters.tint);
  }

  gpu_display_shader_parameters_update(display_shader, display_parameters, float4x4(matrix));
  GPU_uniformbuf_bind(display_shader.parameters_buffer, UniformBufferSlot::DISPLAY);

  /* TODO(fclem): remove remains of IMM. */
  immBindShader(display_shader.shader);

  return true;
}

/** \} */

}  // namespace

GPUShaderBinder::GPUShaderBinder(const Config &config)
    : display_cache_(std::make_unique<internal::GPUShaderCache>()),
      scene_linear_cache_(std::make_unique<internal::GPUShaderCache>()),
      config_(config)
{
}

/* Keep private to the translation unit to allow proper destruction of smart pointers to internal
 * data. */
GPUShaderBinder::~GPUShaderBinder() = default;

bool GPUShaderBinder::display_bind(const GPUDisplayParameters &display_parameters) const
{
  /* Attempt to get shader from the cache. */
  internal::GPUDisplayShader *display_shader = display_cache_->get(display_parameters);

  if (!display_shader) {
    display_shader = &display_cache_->create_default();
    BLI_assert(display_shader);

    display_shader->from_colorspace = display_parameters.from_colorspace;
    display_shader->view = display_parameters.view;
    display_shader->display = display_parameters.display;
    display_shader->look = display_parameters.look;
    display_shader->use_curve_mapping = (display_parameters.curve_mapping != nullptr);
    display_shader->is_valid = false;

    if (display_parameters.curve_mapping) {
      /* Rasterize curve mapping early so that texture allocation can know the size of the lookup
       * table. */
      display_shader->curve_mapping.rasterize(*display_parameters.curve_mapping);
    }

    if (display_shader->initialize_common()) {
      construct_display_shader(*display_shader);
    }
  }
  else if (display_parameters.curve_mapping) {
    /* Update curve mapping's lookup table. */
    display_shader->curve_mapping.rasterize(*display_parameters.curve_mapping);
  }

  return gpu_shader_bind(config_, *display_shader, display_parameters);
}

bool GPUShaderBinder::to_scene_linear_bind(const StringRefNull from_colorspace,
                                           const bool use_predivide) const
{
  /* Re-use code and logic with the conversion to the display space. This assumes that empty names
   * for display, view, and look are not valid for the OpenColorIO configuration, and so they can
   * be used to indicate that the processor is used to convert from the given space to the linear.
   */

  GPUDisplayParameters display_parameters;
  display_parameters.from_colorspace = from_colorspace;
  display_parameters.use_predivide = use_predivide;

  /* Attempt to get shader from the cache. */
  internal::GPUDisplayShader *display_shader = scene_linear_cache_->get(display_parameters);

  if (!display_shader) {
    display_shader = &display_cache_->create_default();
    BLI_assert(display_shader);

    display_shader->from_colorspace = display_parameters.from_colorspace;

    if (display_shader->initialize_common()) {
      construct_scene_linear_shader(*display_shader);
    }
  }

  return gpu_shader_bind(config_, *display_shader, display_parameters);
}

void GPUShaderBinder::unbind() const
{
  immUnbindProgram();
}

bool GPUShaderBinder::create_gpu_shader(
    internal::GPUDisplayShader &display_shader,
    StringRefNull fragment_source,
    const Span<std::array<StringRefNull, 2>> additional_defines)
{
  using namespace blender::gpu::shader;

  StageInterfaceInfo iface("OCIO_Interface", "");
  iface.smooth(Type::float2_t, "texCoord_interp");

  ShaderCreateInfo info("OCIO_Display");

  for (const auto &additional_define : additional_defines) {
    info.define(additional_define[0], additional_define[1]);
  }

  /* Work around OpenColorIO not supporting latest GLSL yet. */
  info.define("texture1D", "texture");
  info.define("texture2D", "texture");
  info.define("texture3D", "texture");

  /* Work around unsupported in keyword in Metal GLSL emulation. */
#if OS_MAC
  info.define("in", "");
#endif

  info.typedef_source("ocio_shader_shared.hh");
  info.sampler(internal::TextureSlot::IMAGE, ImageType::Float2D, "image_texture");
  info.sampler(internal::TextureSlot::OVERLAY, ImageType::Float2D, "overlay_texture");
  info.uniform_buf(internal::UniformBufferSlot::DISPLAY, "OCIO_GPUParameters", "parameters");
  info.push_constant(Type::float4x4_t, "ModelViewProjectionMatrix");
  info.vertex_in(0, Type::float2_t, "pos");
  info.vertex_in(1, Type::float2_t, "texCoord");
  info.vertex_out(iface);
  info.fragment_out(0, Type::float4_t, "fragColor");
  info.vertex_source("gpu_shader_display_transform_vert.glsl");
  info.fragment_source("gpu_shader_display_transform_frag.glsl");

  info.fragment_source_generated = fragment_source;
  process_source(info.fragment_source_generated);

  /* #96502: Work around for incorrect OCIO GLSL code generation when using
   * GradingPrimaryTransform. Should be reevaluated when changing to a next version of OCIO.
   * (currently v2.1.1). */
  info.define("inf 1e32");

  if (display_shader.use_curve_mapping) {
    info.define("USE_CURVE_MAPPING");
    info.uniform_buf(internal::UniformBufferSlot::CURVE_MAPPING,
                     "OCIO_GPUCurveMappingParameters",
                     "curve_mapping");
    info.sampler(
        internal::TextureSlot::CURVE_MAPPING, ImageType::Float1D, "curve_mapping_texture");
  }

  /* Set LUT textures. */
  int slot = internal::TextureSlot::LUTS_OFFSET;
  for (const internal::GPULutTexture &texture : display_shader.textures.luts) {
    const int dimensions = GPU_texture_dimensions(texture.texture);
    const ImageType type = (dimensions == 1) ? ImageType::Float1D :
                           (dimensions == 2) ? ImageType::Float2D :
                                               ImageType::Float3D;
    info.sampler(slot++, type, texture.sampler_name.c_str());
  }

  /* Set LUT uniforms. */
#if defined(WITH_OPENCOLORIO)
  if (!display_shader.textures.uniforms.is_empty()) {
    /* NOTE: For simplicity, we pad everything to size of vec4 avoiding sorting and alignment
     * issues. It is unlikely that this becomes a real issue. */
    const size_t ubo_size = display_shader.textures.uniforms.size() * sizeof(float) * 4;
    Vector<uint8_t> ubo_data_buf(ubo_size);

    uint32_t *ubo_data = reinterpret_cast<uint32_t *>(ubo_data_buf.data());

    std::stringstream ss;
    ss << "struct OCIO_GPULutParameters {\n";

    int index = 0;
    for (internal::GPUUniform &uniform : display_shader.textures.uniforms) {
      index += 1;
      const OCIO_NAMESPACE::GpuShaderDesc::UniformData &data = uniform.data;
      const char *name = uniform.name.c_str();
      char prefix = ' ';
      int vec_len;
      switch (data.m_type) {
        case OCIO_NAMESPACE::UNIFORM_DOUBLE: {
          vec_len = 1;
          float value = float(data.m_getDouble());
          memcpy(ubo_data, &value, sizeof(float));
          break;
        }
        case OCIO_NAMESPACE::UNIFORM_BOOL: {
          prefix = 'b';
          vec_len = 1;
          int value = int(data.m_getBool());
          memcpy(ubo_data, &value, sizeof(int));
          break;
        }
        case OCIO_NAMESPACE::UNIFORM_FLOAT3:
          vec_len = 3;
          memcpy(ubo_data, data.m_getFloat3().data(), sizeof(float) * 3);
          break;
        case OCIO_NAMESPACE::UNIFORM_VECTOR_FLOAT:
          vec_len = data.m_vectorFloat.m_getSize();
          memcpy(ubo_data, data.m_vectorFloat.m_getVector(), sizeof(float) * vec_len);
          break;
        case OCIO_NAMESPACE::UNIFORM_VECTOR_INT:
          prefix = 'i';
          vec_len = data.m_vectorInt.m_getSize();
          memcpy(ubo_data, data.m_vectorInt.m_getVector(), sizeof(int) * vec_len);
          break;
        default:
          continue;
      }
      /* Align every member to 16 bytes. */
      ubo_data += 4;
      /* Use a generic variable name because some GLSL compilers can interpret the preprocessor
       * define as recursive. */
      ss << "  " << prefix << "vec4 var" << index << ";\n";
      /* Use a define to keep the generated code working. */
      StringRef suffix = StringRefNull("xyzw").substr(0, vec_len);
      ss << "#define " << name << " lut_parameters.var" << index << "." << suffix << "\n";
    }
    ss << "};\n";
    info.typedef_source_generated = ss.str();

    info.uniform_buf(internal::UniformBufferSlot::LUTS, "OCIO_GPULutParameters", "lut_parameters");

    display_shader.textures.uniforms_buffer = GPU_uniformbuf_create_ex(
        ubo_size, ubo_data_buf.data(), "OCIO_LutParameters");
  }
#endif

  display_shader.shader = GPU_shader_create_from_info(
      reinterpret_cast<GPUShaderCreateInfo *>(&info));

  return (display_shader.shader != nullptr);
}

}  // namespace blender::ocio
