#include "shader_program.h"

#include "gl_util.h"
#include "shader.h"
#include "texture.h"
#include <gfx_assert.h>

#include <log/log.h>

#include <stdlib.h>
#include <string.h>

/// Creates an OpenGL shader program.
/// Returns non-zero on success, 0 on failure.
static GLuint create_program(GLuint vertex_shader, GLuint fragment_shader) {
  const GLuint prog = glCreateProgram();
  if (prog == 0) {
    LOGE("Failed creating shader program");
    return 0;
  }
  glAttachShader(prog, vertex_shader);
  glAttachShader(prog, fragment_shader);
  glLinkProgram(prog);
  GLint result;
  glGetProgramiv(prog, GL_LINK_STATUS, &result);
  if (result == GL_FALSE) {
    GLint log_len;
    glGetProgramiv(prog, GL_INFO_LOG_LENGTH, &log_len);
    if (log_len > 0) {
      char* log = calloc(log_len, sizeof(char));
      glGetProgramInfoLog(prog, log_len, NULL, log);
      LOGE("Failed creating shader program: %s", log);
      free(log);
    } else {
      LOGE("Failed creating shader program");
    }
    glDeleteProgram(prog);
    return 0;
  }
  ASSERT_GL;
  return prog;
}

bool gfx_build_shader_program(
    ShaderProgram* prog, const ShaderProgramDesc* desc) {
  assert(prog);
  assert(desc);

  prog->id = create_program(desc->vertex_shader->id, desc->fragment_shader->id);
  return prog->id != 0;
}

void gfx_del_shader_program(ShaderProgram* prog) {
  assert(prog);

  if (prog->id) {
    glDeleteProgram(prog->id);
    prog->id = 0;
  }
  ASSERT_GL;
}

void gfx_activate_shader_program(const ShaderProgram* prog) {
  assert(prog);
  glUseProgram(prog->id);
  ASSERT_GL;
}

void gfx_deactivate_shader_program(const ShaderProgram* prog) {
  assert(prog);
  glUseProgram(0);
  ASSERT_GL;
}

static void set_texture_uniform(
    GLuint prog, const char* name, int texture_unit, const Texture* texture) {
  assert(prog != 0);
  assert(name);
  assert(texture);

  const GLint location = glGetUniformLocation(prog, name);
  if (location >= 0) {
    glActiveTexture(GL_TEXTURE0 + texture_unit);
    glBindTexture(texture->target, texture->id);
    glUniform1i(location, texture_unit);
  }
}

static void set_mat4_uniform(
    GLuint prog, const char* name, const mat4* mats, size_t count) {
  assert(prog != 0);
  assert(name);
  assert(mats);

  const GLint location = glGetUniformLocation(prog, name);
  if (location >= 0) {
    glUniformMatrix4fv(location, count, GL_FALSE, (const float*)mats);
  }
}

static void set_vec3_uniform(GLuint prog, const char* name, vec3 value) {
  assert(prog != 0);
  assert(name);

  const GLint location = glGetUniformLocation(prog, name);
  if (location >= 0) {
    glUniform3f(location, value.x, value.y, value.z);
  }
}

static void set_vec4_uniform(GLuint prog, const char* name, vec4 value) {
  assert(prog != 0);
  assert(name);

  const GLint location = glGetUniformLocation(prog, name);
  if (location >= 0) {
    glUniform4f(location, value.x, value.y, value.z, value.w);
  }
}

static void set_float_uniform(GLuint prog, const char* name, float value) {
  assert(prog != 0);
  assert(name);

  const GLint location = glGetUniformLocation(prog, name);
  if (location >= 0) {
    glUniform1f(location, value);
  }
}

void gfx_apply_uniforms(const ShaderProgram* prog) {
  assert(prog);

  int next_texture_unit = 0;
  for (int i = 0; i < prog->num_uniforms; ++i) {
    const ShaderUniform* uniform = &prog->uniforms[i];
    switch (uniform->type) {
    case UniformTexture:
      set_texture_uniform(
          prog->id, uniform->name.str, next_texture_unit,
          uniform->value.texture);
      next_texture_unit++;
      break;
    case UniformMat4:
      set_mat4_uniform(prog->id, uniform->name.str, &uniform->value.mat4, 1);
      break;
    case UniformVec3:
      set_vec3_uniform(prog->id, uniform->name.str, uniform->value.vec3);
      break;
    case UniformVec4:
      set_vec4_uniform(prog->id, uniform->name.str, uniform->value.vec4);
      break;
    case UniformFloat:
      set_float_uniform(prog->id, uniform->name.str, uniform->value.scalar);
      break;
    case UniformMat4Array:
      set_mat4_uniform(
          prog->id, uniform->name.str, uniform->value.array.values,
          uniform->value.array.count);
      break;
    }
  }
}

// Get the ShaderUniform object by name from the shader program if it already
// exists, or allocate a new one otherwise.
static ShaderUniform* get_or_allocate_uniform(
    ShaderProgram* prog, const char* name) {
  assert(prog);
  assert(name);

  // First search for the uniform in the list.
  for (int i = 0; i < prog->num_uniforms; ++i) {
    ShaderUniform* uniform = &prog->uniforms[i];
    if (sstring_eq_cstr(uniform->name, name)) {
      return uniform;
    }
  }

  // Create the uniform if it does not exist.
  if (prog->num_uniforms == GFX_MAX_UNIFORMS_PER_SHADER) {
    FAIL("Exceeded the maximum number of uniforms per shader. Please increase "
         "this value.");
    return 0;
  }
  ShaderUniform* uniform = &prog->uniforms[prog->num_uniforms];
  prog->num_uniforms++;
  return uniform;
}

// The functions below save the value of a uniform in the shader program. If the
// uniform does not even exist, then there is no need to store the value.

void gfx_set_texture_uniform(
    ShaderProgram* prog, const char* name, const Texture* texture) {
  assert(prog);
  assert(name);
  assert(texture);

  const GLint location = glGetUniformLocation(prog->id, name);
  if (location < 0) {
    return;
  }
  ShaderUniform* uniform = get_or_allocate_uniform(prog, name);
  assert(uniform);
  uniform->name          = sstring_make(name);
  uniform->type          = UniformTexture;
  uniform->value.texture = texture;
}

void gfx_set_mat4_uniform(
    ShaderProgram* prog, const char* name, const mat4* mat) {
  assert(prog);
  assert(name);
  assert(mat);

  const GLint location = glGetUniformLocation(prog->id, name);
  if (location < 0) {
    return;
  }
  ShaderUniform* uniform = get_or_allocate_uniform(prog, name);
  assert(uniform);
  uniform->name       = sstring_make(name);
  uniform->type       = UniformMat4;
  uniform->value.mat4 = *mat;
}

void gfx_set_vec3_uniform(ShaderProgram* prog, const char* name, vec3 value) {
  assert(prog);
  assert(name);

  const GLint location = glGetUniformLocation(prog->id, name);
  if (location < 0) {
    return;
  }
  ShaderUniform* uniform = get_or_allocate_uniform(prog, name);
  assert(uniform);
  uniform->name       = sstring_make(name);
  uniform->type       = UniformVec3;
  uniform->value.vec3 = value;
}

void gfx_set_vec4_uniform(ShaderProgram* prog, const char* name, vec4 value) {
  assert(prog);
  assert(name);

  const GLint location = glGetUniformLocation(prog->id, name);
  if (location < 0) {
    return;
  }
  ShaderUniform* uniform = get_or_allocate_uniform(prog, name);
  assert(uniform);
  uniform->name       = sstring_make(name);
  uniform->type       = UniformVec4;
  uniform->value.vec4 = value;
}

void gfx_set_float_uniform(ShaderProgram* prog, const char* name, float value) {
  assert(prog);
  assert(name);

  // No need to store the uniform on our side if it does not exist in the
  // program.
  const GLint location = glGetUniformLocation(prog->id, name);
  if (location < 0) {
    return;
  }
  ShaderUniform* uniform = get_or_allocate_uniform(prog, name);
  assert(uniform);
  uniform->name         = sstring_make(name);
  uniform->type         = UniformFloat;
  uniform->value.scalar = value;
}

void gfx_set_mat4_array_uniform(
    ShaderProgram* prog, const char* name, const mat4* mats, size_t count) {
  assert(prog);
  assert(name);
  assert(mats);
  
  const GLint location = glGetUniformLocation(prog->id, name);
  if (location < 0) {
    return;
  }
  ShaderUniform* uniform = get_or_allocate_uniform(prog, name);
  assert(uniform);
  uniform->name               = sstring_make(name);
  uniform->type               = UniformMat4Array;
  uniform->value.array.count  = count;
  uniform->value.array.values = mats;
}