/******************************************************************************
  This source file is part of the Avogadro project.
  This source code is released under the 3-Clause BSD License, (see "LICENSE").
******************************************************************************/

#include "shaderprogram.h"

#include "avogadrogl.h"
#include "shader.h"
#include "texture2d.h"

#include <avogadro/core/types.h>

#include <algorithm>

namespace Avogadro::Rendering {

namespace {
inline GLenum convertType(Type type)
{
  switch (type) {
    default:
    case UCharType:
      return GL_UNSIGNED_BYTE;
    case CharType:
      return GL_BYTE;
    case ShortType:
      return GL_SHORT;
    case UShortType:
      return GL_UNSIGNED_SHORT;
    case IntType:
      return GL_INT;
    case UIntType:
      return GL_UNSIGNED_INT;
    case FloatType:
      return GL_FLOAT;
    case DoubleType:
      return GL_DOUBLE;
  }
}

inline GLenum lookupTextureUnit(GLint index)
{
#define MAKE_TEXTURE_UNIT_CASE(i)                                              \
  case i:                                                                      \
    return GL_TEXTURE##i;
  switch (index) {
    MAKE_TEXTURE_UNIT_CASE(0)
    MAKE_TEXTURE_UNIT_CASE(1)
    MAKE_TEXTURE_UNIT_CASE(2)
    MAKE_TEXTURE_UNIT_CASE(3)
    MAKE_TEXTURE_UNIT_CASE(4)
    MAKE_TEXTURE_UNIT_CASE(5)
    MAKE_TEXTURE_UNIT_CASE(6)
    MAKE_TEXTURE_UNIT_CASE(7)
    MAKE_TEXTURE_UNIT_CASE(8)
    MAKE_TEXTURE_UNIT_CASE(9)
    MAKE_TEXTURE_UNIT_CASE(10)
    MAKE_TEXTURE_UNIT_CASE(11)
    MAKE_TEXTURE_UNIT_CASE(12)
    MAKE_TEXTURE_UNIT_CASE(13)
    MAKE_TEXTURE_UNIT_CASE(14)
    MAKE_TEXTURE_UNIT_CASE(15)
    MAKE_TEXTURE_UNIT_CASE(16)
    MAKE_TEXTURE_UNIT_CASE(17)
    MAKE_TEXTURE_UNIT_CASE(18)
    MAKE_TEXTURE_UNIT_CASE(19)
    MAKE_TEXTURE_UNIT_CASE(20)
    MAKE_TEXTURE_UNIT_CASE(21)
    MAKE_TEXTURE_UNIT_CASE(22)
    MAKE_TEXTURE_UNIT_CASE(23)
    MAKE_TEXTURE_UNIT_CASE(24)
    MAKE_TEXTURE_UNIT_CASE(25)
    MAKE_TEXTURE_UNIT_CASE(26)
    MAKE_TEXTURE_UNIT_CASE(27)
    MAKE_TEXTURE_UNIT_CASE(28)
    MAKE_TEXTURE_UNIT_CASE(29)
    MAKE_TEXTURE_UNIT_CASE(30)
    MAKE_TEXTURE_UNIT_CASE(31)
    default:
      return 0;
  }
}
} // namespace

ShaderProgram::ShaderProgram()
  : m_handle(0), m_vertexShader(0), m_fragmentShader(0), m_linked(false)
{
  initializeTextureUnits();
}

ShaderProgram::~ShaderProgram()
{
  if (m_handle != 0) {
    glDeleteProgram(static_cast<GLuint>(m_handle));
  }
}

bool ShaderProgram::attachShader(const Shader& shader)
{
  if (shader.handle() == 0) {
    m_error = "Shader object was not initialized, cannot attach it.";
    return false;
  }
  if (shader.type() == Shader::Unknown) {
    m_error = "Shader object is of type Unknown and cannot be used.";
    return false;
  }

  if (m_handle == 0) {
    GLuint handle_ = glCreateProgram();
    if (handle_ == 0) {
      m_error = "Could not create shader program.";
      return false;
    }
    m_handle = static_cast<Index>(handle_);
    m_linked = false;
  }

  if (shader.type() == Shader::Vertex) {
    if (m_vertexShader != 0) {
      glDetachShader(static_cast<GLuint>(m_handle),
                     static_cast<GLuint>(m_vertexShader));
    }
    m_vertexShader = shader.handle();
  } else if (shader.type() == Shader::Fragment) {
    if (m_fragmentShader != 0) {
      glDetachShader(static_cast<GLuint>(m_handle),
                     static_cast<GLuint>(m_fragmentShader));
    }
    m_fragmentShader = shader.handle();
  } else {
    m_error = "Unknown shader type encountered - this should not happen.";
    return false;
  }

  glAttachShader(static_cast<GLuint>(m_handle),
                 static_cast<GLuint>(shader.handle()));
  m_linked = false;
  return true;
}

bool ShaderProgram::detachShader(const Shader& shader)
{
  if (shader.handle() == 0) {
    m_error = "Shader object was not initialized, cannot attach it.";
    return false;
  }
  if (shader.type() == Shader::Unknown) {
    m_error = "Shader object is of type Unknown and cannot be used.";
    return false;
  }
  if (m_handle == 0) {
    m_error = "This shader prorgram has not been initialized yet.";
  }

  switch (shader.type()) {
    case Shader::Vertex:
      if (m_vertexShader != shader.handle()) {
        m_error = "The supplied shader was not attached to this program.";
        return false;
      } else {
        glDetachShader(static_cast<GLuint>(m_handle),
                       static_cast<GLuint>(shader.handle()));
        m_vertexShader = 0;
        return true;
      }
    case Shader::Fragment:
      if (m_fragmentShader != shader.handle()) {
        m_error = "The supplied shader was not attached to this program.";
        return false;
      } else {
        glDetachShader(static_cast<GLuint>(m_handle),
                       static_cast<GLuint>(shader.handle()));
        m_fragmentShader = 0;
        return true;
      }
    default:
      return false;
  }
}

bool ShaderProgram::link()
{
  if (m_linked)
    return true;

  if (m_handle == 0) {
    m_error = "Program has not been initialized, and/or does not have shaders.";
    return false;
  }

  GLint isCompiled;
  glLinkProgram(static_cast<GLuint>(m_handle));
  glGetProgramiv(static_cast<GLuint>(m_handle), GL_LINK_STATUS, &isCompiled);
  if (isCompiled == 0) {
    GLint length(0);
    glGetShaderiv(static_cast<GLuint>(m_handle), GL_INFO_LOG_LENGTH, &length);
    if (length > 1) {
      char* logMessage = new char[length];
      glGetShaderInfoLog(static_cast<GLuint>(m_handle), length, nullptr,
                         logMessage);
      m_error = logMessage;
      delete[] logMessage;
    }
    return false;
  }
  m_linked = true;
  m_attributes.clear();
  return true;
}

bool ShaderProgram::bind()
{
  if (!m_linked && !link())
    return false;

  glUseProgram(static_cast<GLuint>(m_handle));
  return true;
}

void ShaderProgram::release()
{
  glUseProgram(0);
  releaseAllTextureUnits();
}

bool ShaderProgram::enableAttributeArray(const std::string& name)
{
  auto location = static_cast<GLint>(findAttributeArray(name));
  if (location == -1) {
    m_error = "Could not enable attribute " + name + ". No such attribute.";
    return false;
  }
  glEnableVertexAttribArray(location);
  return true;
}

bool ShaderProgram::disableAttributeArray(const std::string& name)
{
  auto location = static_cast<GLint>(findAttributeArray(name));
  if (location == -1) {
    m_error = "Could not disable attribute " + name + ". No such attribute.";
    return false;
  }
  glDisableVertexAttribArray(location);
  return true;
}

#define BUFFER_OFFSET(i) ((char*)nullptr + (i))

bool ShaderProgram::useAttributeArray(const std::string& name, int offset,
                                      size_t stride, Type elementType,
                                      int elementTupleSize,
                                      NormalizeOption normalize)
{
  auto location = static_cast<GLint>(findAttributeArray(name));
  if (location == -1) {
    m_error = "Could not use attribute " + name + ". No such attribute.";
    return false;
  }
  glVertexAttribPointer(location, elementTupleSize, convertType(elementType),
                        normalize == Normalize ? GL_TRUE : GL_FALSE,
                        static_cast<GLsizei>(stride), BUFFER_OFFSET(offset));
  return true;
}

bool ShaderProgram::setTextureSampler(const std::string& name,
                                      const Texture2D& texture)
{
  // Look up sampler location:
  auto location = static_cast<GLint>(findUniform(name));
  if (location == -1) {
    m_error = "Could not set sampler " + name + ". No uniform with that name.";
    return false;
  }

  // Check if the texture is already bound:
  GLint textureUnitId = 0;
  auto result = m_textureUnitBindings.find(&texture);
  if (result == m_textureUnitBindings.end()) {
    // Not bound. Attempt to bind the texture to an available texture unit.
    // We'll leave GL_TEXTURE0 unbound, as it is used for manipulating
    // textures.
    auto begin = m_boundTextureUnits.begin() + 1;
    auto end = m_boundTextureUnits.end();
    auto available = std::find(begin, end, false);

    if (available == end) {
      m_error = "Could not set sampler " + name +
                ". No remaining texture units available.";
      return false;
    }

    textureUnitId = static_cast<GLint>(available - begin);

    GLenum textureUnit = lookupTextureUnit(textureUnitId);
    if (textureUnit == 0) {
      m_error =
        "Could not set sampler " + name + ". Texture unit lookup failed.";
      return false;
    }

    glActiveTexture(textureUnit);
    if (!texture.bind()) {
      m_error = "Could not set sampler " + name +
                ": Error while binding texture: '" + texture.error() + "'.";
      glActiveTexture(GL_TEXTURE0);
      return false;
    }
    glActiveTexture(GL_TEXTURE0);

    // Mark texture unit as in-use.
    m_textureUnitBindings.insert(std::make_pair(&texture, textureUnitId));
    *available = true;
  } else {
    // Texture is already bound.
    textureUnitId = result->second;
  }

  // Set the texture unit uniform
  glUniform1i(location, textureUnitId);

  return true;
}

bool ShaderProgram::setUniformValue(const std::string& name, int i)
{
  auto location = static_cast<GLint>(findUniform(name));
  if (location == -1) {
    m_error = "Could not set uniform " + name + ". No such uniform.";
    return false;
  }
  glUniform1i(location, static_cast<GLint>(i));
  return true;
}

bool ShaderProgram::setUniformValue(const std::string& name, float f)
{
  auto location = static_cast<GLint>(findUniform(name));
  if (location == -1) {
    m_error = "Could not set uniform " + name + ". No such uniform.";
    return false;
  }
  glUniform1f(location, static_cast<GLfloat>(f));
  return true;
}

bool ShaderProgram::setUniformValue(const std::string& name,
                                    const Eigen::Matrix3f& matrix)
{
  auto location = static_cast<GLint>(findUniform(name));
  if (location == -1) {
    m_error = "Could not set uniform " + name + ". No such uniform.";
    return false;
  }
  glUniformMatrix3fv(location, 1, GL_FALSE,
                     static_cast<const GLfloat*>(matrix.data()));
  return true;
}

bool ShaderProgram::setUniformValue(const std::string& name,
                                    const Eigen::Matrix4f& matrix)
{
  auto location = static_cast<GLint>(findUniform(name));
  if (location == -1) {
    m_error = "Could not set uniform " + name + ". No such uniform.";
    return false;
  }
  glUniformMatrix4fv(location, 1, GL_FALSE,
                     static_cast<const GLfloat*>(matrix.data()));
  return true;
}

bool ShaderProgram::setUniformValue(const std::string& name, const Vector3f& v)
{
  auto location = static_cast<GLint>(findUniform(name));
  if (location == -1) {
    m_error = "Could not set uniform " + name + ". No such uniform.";
    return false;
  }
  glUniform3fv(location, 1, v.data());
  return true;
}

bool ShaderProgram::setUniformValue(const std::string& name, const Vector2i& v)
{
  auto location = static_cast<GLint>(findUniform(name));
  if (location == -1) {
    m_error = "Could not set uniform " + name + ". No such uniform.";
    return false;
  }
  glUniform2iv(location, 1, v.data());
  return true;
}

bool ShaderProgram::setUniformValue(const std::string& name, const Vector3ub& v)
{
  auto location = static_cast<GLint>(findUniform(name));
  if (location == -1) {
    m_error = "Could not set uniform " + name + ". No such uniform.";
    return false;
  }
  Vector3f colorf(v.cast<float>() * (1.0f / 255.0f));
  glUniform3fv(location, 1, colorf.data());
  return true;
}

bool ShaderProgram::setAttributeArrayInternal(
  const std::string& name, void* buffer, Avogadro::Type type, int tupleSize,
  ShaderProgram::NormalizeOption normalize)
{
  if (type == Avogadro::UnknownType) {
    m_error = "Unrecognized data type for attribute " + name + ".";
    return false;
  }
  auto location = static_cast<GLint>(findAttributeArray(name));
  if (location == -1) {
    m_error = "Could not set attribute " + name + ". No such attribute.";
    return false;
  }
  const auto* data = static_cast<const GLvoid*>(buffer);
  glVertexAttribPointer(location, tupleSize, convertType(type),
                        normalize == Normalize ? GL_TRUE : GL_FALSE, 0, data);
  return true;
}

void ShaderProgram::initializeTextureUnits()
{
  GLint numTextureUnits = 0;
  glGetIntegerv(GL_MAX_COMBINED_TEXTURE_IMAGE_UNITS, &numTextureUnits);

  // We'll impose a hard limit of 32 texture units for symbolic lookups.
  // This seems to be about the maximum available on current hardware.
  // If increasing this limit, modify the lookupTextureUnit method
  // appropriately.
  numTextureUnits = std::clamp(numTextureUnits, 0, 32);

  m_boundTextureUnits.clear();
  m_boundTextureUnits.resize(numTextureUnits, false);
  m_textureUnitBindings.clear();
}

void ShaderProgram::releaseAllTextureUnits()
{
  std::fill(m_boundTextureUnits.begin(), m_boundTextureUnits.end(), false);
  m_textureUnitBindings.clear();
}

inline int ShaderProgram::findAttributeArray(const std::string& name)
{
  if (name.empty() || !m_linked)
    return -1;
  const auto* namePtr = static_cast<const GLchar*>(name.c_str());
  auto location = static_cast<int>(
    glGetAttribLocation(static_cast<GLuint>(m_handle), namePtr));
  if (location == -1) {
    m_error = "Specified attribute not found in current shader program: ";
    m_error += name;
  }

  return location;
}

inline int ShaderProgram::findUniform(const std::string& name)
{
  if (name.empty() || !m_linked)
    return -1;
  const auto* namePtr = static_cast<const GLchar*>(name.c_str());
  auto location = static_cast<int>(
    glGetUniformLocation(static_cast<GLuint>(m_handle), namePtr));
  if (location == -1)
    m_error = "Uniform " + name + " not found in current shader program.";

  return location;
}

} // namespace Avogadro::Rendering
