/******************************************************************************

  This source file is part of the Avogadro project.

  Copyright 2012 Kitware, Inc.

  This source code is released under the New BSD License, (the "License").

  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.

******************************************************************************/

#include "shaderprogram.h"

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

#include <avogadro/core/types.h>

#include <algorithm>

namespace Avogadro {
namespace 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;
  }
}
} // end anon namespace

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

ShaderProgram::~ShaderProgram()
{
}

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;
        m_linked = false;
        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;
        m_linked = false;
        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)
{
  GLint 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)
{
  GLint 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)
{
  GLint 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:
  GLint 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;
  typedef std::map<const Texture2D*, int>::const_iterator TMapIter;
  TMapIter 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.
    std::vector<bool>::iterator begin = m_boundTextureUnits.begin() + 1;
    std::vector<bool>::iterator end = m_boundTextureUnits.end();
    std::vector<bool>::iterator 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)
{
  GLint 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)
{
  GLint 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)
{
  GLint 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)
{
  GLint 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)
{
  GLint 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)
{
  GLint 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)
{
  GLint 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;
  }
  GLint location = static_cast<GLint>(findAttributeArray(name));
  if (location == -1) {
    m_error = "Could not set attribute " + name + ". No such attribute.";
    return false;
  }
  const GLvoid* 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;
  glGetIntegerv(GL_MAX_COMBINED_TEXTURE_IMAGE_UNITS, &numTextureUnits);

  // We'll impose a 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::min(std::max(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 GLchar* namePtr = static_cast<const GLchar*>(name.c_str());
  GLint 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 GLchar* namePtr = static_cast<const GLchar*>(name.c_str());
  GLint 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;
}

} // End Rendering namespace
} // End Avogadro namespace
