#include "imm_renderer_impl.h"
#include "light_impl.h"
#include "mesh_impl.h"

#include "scene/animation_impl.h"

#include <gfx/core.h>
#include <gfx/util/ibl.h>
#include <gfx/util/shader.h>

#include <math/aabb3.h>

#include <cassert.h>
#include <string.h> // memcpy

static const int IRRADIANCE_MAP_WIDTH               = 1024;
static const int IRRADIANCE_MAP_HEIGHT              = 1024;
static const int PREFILTERED_ENVIRONMENT_MAP_WIDTH  = 128;
static const int PREFILTERED_ENVIRONMENT_MAP_HEIGHT = 128;
static const int BRDF_INTEGRATION_MAP_WIDTH         = 512;
static const int BRDF_INTEGRATION_MAP_HEIGHT        = 512;

/// Initialize renderer state for IBL.
static bool init_ibl(ImmRenderer* renderer) {
  assert(renderer);
  assert(!renderer->ibl);
  assert(!renderer->brdf_integration_map);

  if (!((renderer->ibl = gfx_make_ibl(renderer->gfxcore)))) {
    return false;
  }

  if (!((renderer->brdf_integration_map = gfx_make_brdf_integration_map(
             renderer->ibl, renderer->gfxcore, BRDF_INTEGRATION_MAP_WIDTH,
             BRDF_INTEGRATION_MAP_HEIGHT)))) {
    return false;
  }

  return true;
}

// TODO: Why is this done lazily here? Do it when the environment light is
//  created.
//
/// Compute irradiance and prefiltered environment maps for the light if they
/// have not been already computed.
static bool set_up_environment_light(
    ImmRenderer* renderer, EnvironmentLight* light) {
  assert(renderer);
  assert(light);
  assert(renderer->ibl);
  assert(renderer->brdf_integration_map);

  if (light->irradiance_map) {
    assert(light->prefiltered_environment_map);
    return true;
  }

  // For convenience.
  GfxCore* gfxcore = renderer->gfxcore;

  Texture* irradiance_map              = 0;
  Texture* prefiltered_environment_map = 0;

  if (!((irradiance_map = gfx_make_irradiance_map(
             renderer->ibl, gfxcore, light->environment_map,
             IRRADIANCE_MAP_WIDTH, IRRADIANCE_MAP_HEIGHT)))) {
    goto cleanup;
  }

  int max_mip_level = 0;
  if (!((prefiltered_environment_map = gfx_make_prefiltered_environment_map(
             renderer->ibl, gfxcore, light->environment_map,
             PREFILTERED_ENVIRONMENT_MAP_WIDTH,
             PREFILTERED_ENVIRONMENT_MAP_HEIGHT, &max_mip_level)))) {
    goto cleanup;
  }

  light->irradiance_map              = irradiance_map;
  light->prefiltered_environment_map = prefiltered_environment_map;
  light->max_reflection_lod          = max_mip_level;

  return true;

cleanup:
  if (irradiance_map) {
    gfx_destroy_texture(gfxcore, &irradiance_map);
  }
  if (prefiltered_environment_map) {
    gfx_destroy_texture(gfxcore, &prefiltered_environment_map);
  }
  return false;
}

static void configure_light(ImmRenderer* renderer, Light* light) {
  assert(renderer);
  assert(light);

  // For convenience.
  ShaderProgram* const shader = renderer->shader;

  switch (light->type) {
  case EnvironmentLightType: {
    EnvironmentLight* env = &light->environment;

    const bool initialized = set_up_environment_light(renderer, env);
    ASSERT(initialized);
    assert(env->environment_map);
    assert(env->irradiance_map);
    assert(env->prefiltered_environment_map);
    assert(renderer->brdf_integration_map);

    gfx_set_texture_uniform(
        shader, "BRDFIntegrationMap", renderer->brdf_integration_map);
    gfx_set_texture_uniform(shader, "Sky", env->environment_map);
    gfx_set_texture_uniform(shader, "IrradianceMap", env->irradiance_map);
    gfx_set_texture_uniform(
        shader, "PrefilteredEnvironmentMap", env->prefiltered_environment_map);
    gfx_set_float_uniform(
        shader, "MaxReflectionLOD", (float)env->max_reflection_lod);

    break;
  }
  default:
    assert(false); // TODO: Implement other light types.
    break;
  }
}

static void configure_state(ImmRenderer* renderer) {
  assert(renderer);

  // Check if anything changed first so that we don't call gfx_apply_uniforms()
  // unnecessarily.
  const bool anything_changed =
      renderer->camera_changed || renderer->lights_changed ||
      renderer->skeleton_changed || renderer->shader_changed;
  if (!anything_changed) {
    return;
  }

  // For convenience.
  ShaderProgram* const shader = renderer->shader;

  // TODO: camera_changed is not set anywhere. Need to think how imm primitive
  //  rendering and imm mesh rendering work together. We could treat imm
  //  primitive calls like setting a new shader.
  if (renderer->camera_changed || renderer->shader_changed) {
    renderer->camera_changed = false;

    // Set all supported camera-related uniforms. Shaders can choose which ones
    // to use.
    // TODO: Check to see which ones the shader actually uses and avoid
    //  computing the unnecessary matrices.
    const mat4* const model = &renderer->matrix_stack[renderer->stack_pointer];
    const mat4        modelview = mat4_mul(renderer->view, *model);
    const mat4 view_proj = mat4_mul(renderer->projection, renderer->view);
    const mat4 mvp       = mat4_mul(renderer->projection, modelview);

    gfx_set_mat4_uniform(shader, "ModelMatrix", model);
    gfx_set_mat4_uniform(shader, "Modelview", &modelview);
    gfx_set_mat4_uniform(shader, "View", &renderer->view);
    gfx_set_mat4_uniform(shader, "Projection", &renderer->projection);
    gfx_set_mat4_uniform(shader, "ViewProjection", &view_proj);
    gfx_set_mat4_uniform(shader, "MVP", &mvp);
    gfx_set_vec3_uniform(shader, "CameraPosition", renderer->camera_position);
  }

  if (renderer->lights_changed || renderer->shader_changed) {
    renderer->lights_changed = false;

    // TODO: Could do better by only setting the lights that have actually
    //  changed.
    // TODO: Will also need to pass the number of lights to the shader once the
    //  other light types are implemented.
    for (int i = 0; i < renderer->num_lights; ++i) {
      configure_light(renderer, renderer->lights[i]);
    }
  }

  if (renderer->skeleton_changed || renderer->shader_changed) {
    renderer->skeleton_changed = false;

    gfx_set_mat4_array_uniform(
        shader, "JointMatrices", renderer->joint_matrices,
        renderer->num_joints);
  }

  if (renderer->shader_changed) {
    renderer->shader_changed = false;
    gfx_activate_shader_program(renderer->shader);
  }

  // Must be called after activating the program.
  gfx_apply_uniforms(renderer->shader);
}

bool gfx_imm_make(ImmRenderer* renderer, GfxCore* gfxcore) {
  assert(renderer);
  assert(gfxcore);

  const size_t num_triangle_verts = IMM_MAX_NUM_TRIANGLES * 3;

  renderer->gfxcore = gfxcore;

  renderer->triangles = gfx_make_geometry(
      gfxcore,
      &(GeometryDesc){.type         = Triangles,
                      .buffer_usage = BufferDynamic,
                      .num_verts    = num_triangle_verts,
                      .positions3d  = (BufferView3d){
                           .size_bytes = num_triangle_verts * sizeof(vec3)}});
  if (!renderer->triangles) {
    goto cleanup;
  }

  renderer->imm_shader = gfx_make_immediate_mode_shader(gfxcore);
  if (!renderer->imm_shader) {
    goto cleanup;
  }
  renderer->shader = renderer->imm_shader;

  if (!init_ibl(renderer)) {
    goto cleanup;
  }

  gfx_imm_load_identity(renderer);
  gfx_imm_set_colour(renderer, vec4_make(0.0f, 0.0f, 0.0f, 1.0f));

  return true;

cleanup:
  gfx_imm_destroy(renderer);
  return false;
}

void gfx_imm_destroy(ImmRenderer* renderer) {
  assert(renderer);
  assert(renderer->gfxcore);

  if (renderer->triangles) {
    gfx_destroy_geometry(renderer->gfxcore, &renderer->triangles);
    // TODO: Could also destroy the geometry's buffers here.
  }

  if (renderer->imm_shader) {
    gfx_destroy_shader_program(renderer->gfxcore, &renderer->imm_shader);
  }

  if (renderer->brdf_integration_map) {
    gfx_destroy_texture(renderer->gfxcore, &renderer->brdf_integration_map);
  }

  // TODO: Do this once the IBL from the scene renderer is gone.
  if (renderer->ibl) {
    // gfx_destroy_ibl(renderer->gfxcore, &renderer->ibl);
  }
}

void gfx_imm_flush(ImmRenderer* renderer) {
  assert(renderer);

  if (renderer->num_triangle_verts > 0) {
    configure_state(renderer);

    gfx_update_geometry(
        renderer->triangles,
        &(GeometryDesc){
            .num_verts   = renderer->num_triangle_verts,
            .positions3d = (BufferView3d){
                                          .data       = renderer->triangle_verts,
                                          .size_bytes = renderer->num_triangle_verts * sizeof(vec3)}
    });

    gfx_apply_uniforms(renderer->shader);
    gfx_render_geometry(renderer->triangles);

    renderer->num_triangle_verts = 0;
  }
}

void gfx_imm_set_shader(ImmRenderer* renderer, ShaderProgram* shader) {
  assert(renderer);
  assert(shader);

  // TODO: It would probably be best to make the imm renderer work in terms of a
  //  new LLR renderer. Otherwise we need to constantly flush stuff everywhere
  //  "just in case". This would still allow the imm to render meshes with
  //  lighting etc. We just need to create an actual Mesh out of the 'triangles'
  //  Geometry that imm currently has. The change would greatly simplify the
  //  implementation of this otherwise coupled LLR-IMM renderer.
  //  Need to decide where to put the matrix stack manipulation. Might be good
  //  to move to the LLR. (Currently, manipulating the stack causes an imm
  //  flush, but that's because we have coupled imm with stack manipulation,
  //  which the new design seems like would address.)
  gfx_imm_flush(renderer);

  // It's important to not set shader_changed unnecessarily, since that would
  // re-trigger the setting of uniforms.
  if (renderer->shader != shader) {
    renderer->shader         = shader;
    renderer->shader_changed = true;
  }
}

void gfx_imm_start(ImmRenderer* renderer) {
  assert(renderer);

  // Shader uniforms are applied lazily.
  // TODO: In the event that gfx_activate_shader_program() activates uniforms
  //  automatically for convenience, call an overload here that doesn't do so.
  gfx_activate_shader_program(renderer->shader);
}

void gfx_imm_end(ImmRenderer* renderer) {
  assert(renderer);

  gfx_imm_flush(renderer);
  gfx_deactivate_shader_program(renderer->shader);

  // TODO: Should we clear all of the render state here as well? At least set
  //  the 'changed' variables to false, for example.
}

void gfx_imm_push_light(ImmRenderer* renderer, Light* light) {
  assert(renderer);
  assert(light);
  assert(renderer->num_lights >= 0);
  ASSERT(renderer->num_lights < IMM_MAX_NUM_LIGHTS);

  renderer->lights[renderer->num_lights++] = light;
  renderer->lights_changed                 = true;
}

void gfx_imm_pop_light(ImmRenderer* renderer) {
  assert(renderer);
  ASSERT(renderer->num_lights > 0);

  renderer->lights[--renderer->num_lights] = 0;
  renderer->lights_changed                 = true;
}

void gfx_imm_set_skeleton(
    ImmRenderer* renderer, const Anima* anima, const Skeleton* skeleton) {
  assert(renderer);
  assert(anima);
  assert(skeleton);
  assert(skeleton->num_joints <= GFX_MAX_NUM_JOINTS);

  for (size_t i = 0; i < skeleton->num_joints; ++i) {
    const joint_idx joint_index = skeleton->joints[i];
    const Joint*    joint       = &anima->joints[joint_index];
    renderer->joint_matrices[i] = joint->joint_matrix;
  }
  renderer->num_joints       = skeleton->num_joints;
  renderer->skeleton_changed = true;
}

void gfx_imm_unset_skeleton(ImmRenderer* renderer) {
  assert(renderer);

  renderer->num_joints       = 0;
  renderer->skeleton_changed = true;
}

void gfx_imm_render_mesh(ImmRenderer* renderer, const Mesh* mesh) {
  assert(renderer);
  assert(mesh);
  assert(mesh->geometry);
  assert(mesh->material);

  configure_state(renderer);
  gfx_render_geometry(mesh->geometry);
}

void gfx_imm_draw_triangles(
    ImmRenderer* renderer, const vec3 verts[], size_t num_triangles) {
  assert(renderer);
  assert(verts);
  const size_t new_verts = num_triangles * 3;
  assert(
      renderer->num_triangle_verts + new_verts < (IMM_MAX_NUM_TRIANGLES * 3));

  memcpy(
      renderer->triangle_verts + renderer->num_triangle_verts, verts,
      new_verts * sizeof(vec3));

  renderer->num_triangle_verts += new_verts;
}

void gfx_imm_draw_triangle(ImmRenderer* renderer, const vec3 verts[3]) {
  gfx_imm_draw_triangles(renderer, verts, 1);
}

void gfx_imm_draw_aabb2(ImmRenderer* renderer, aabb2 box) {
  assert(renderer);

  // clang-format off
  const vec3 verts[4] = {
      vec3_make(box.min.x, box.min.y, 0),    // 3 ---- 2
      vec3_make(box.max.x, box.min.y, 0),    // |      |
      vec3_make(box.max.x, box.max.y, 0),    // |      |
      vec3_make(box.min.x, box.max.y, 0)};   // 0 ---- 1
  // clang-format on

#define tri(i0, i1, i2) verts[i0], verts[i1], verts[i2]
  const vec3 tris[6] = {tri(0, 1, 2), tri(0, 2, 3)};
#undef tri

  gfx_imm_draw_triangles(renderer, tris, 2);
}

void gfx_imm_draw_aabb3(ImmRenderer* renderer, aabb3 box) {
  assert(renderer);

  // clang-format off
  const vec3 vertices[8] = {
      vec3_make(box.min.x, box.min.y, box.max.z), //     7 ----- 6
      vec3_make(box.max.x, box.min.y, box.max.z), //    /       /|
      vec3_make(box.max.x, box.max.y, box.max.z), //   3 ----- 2 |
      vec3_make(box.min.x, box.max.y, box.max.z), //   |       | |
      vec3_make(box.min.x, box.min.y, box.min.z), //   | 4 ----- 5
      vec3_make(box.max.x, box.min.y, box.min.z), //   |/      |/
      vec3_make(box.max.x, box.max.y, box.min.z), //   0 ----- 1
      vec3_make(box.min.x, box.max.y, box.min.z)};
  // clang-format on

  gfx_imm_draw_box3(renderer, vertices);
}

void gfx_imm_draw_box3(ImmRenderer* renderer, const vec3 vertices[8]) {
  assert(renderer);
  assert(vertices);

  //     7 ----- 6
  //    /       /|
  //   3 ----- 2 |
  //   |       | |
  //   | 4 ----- 5
  //   |/      |/
  //   0 ----- 1

#define tri(i0, i1, i2) vertices[i0], vertices[i1], vertices[i2]
  const vec3 tris[36] = {
      // Front.
      tri(0, 1, 2), tri(0, 2, 3),
      // Right.
      tri(1, 5, 6), tri(1, 6, 2),
      // Back.
      tri(5, 4, 7), tri(5, 7, 6),
      // Left.
      tri(4, 0, 03), tri(4, 3, 7),
      // Top.
      tri(3, 2, 6), tri(3, 6, 7),
      // Bottom.
      tri(0, 4, 5), tri(0, 5, 1)};

  gfx_imm_draw_triangles(renderer, tris, 12);
}

void gfx_imm_set_colour(ImmRenderer* renderer, vec4 colour) {
  assert(renderer);
  assert(renderer->shader);

  gfx_imm_flush(renderer);

  gfx_set_vec4_uniform(renderer->shader, "Colour", colour);
}

// Load the top of the matrix stack into the shader.
static void update_shader_model_matrix(ImmRenderer* renderer) {
  assert(renderer);

  gfx_imm_flush(renderer);

  gfx_set_mat4_uniform(
      renderer->shader, "Model",
      &renderer->matrix_stack[renderer->stack_pointer]);
}

void gfx_imm_load_identity(ImmRenderer* renderer) {
  assert(renderer);

  renderer->matrix_stack[0] = mat4_id();
  renderer->stack_pointer   = 0;
  update_shader_model_matrix(renderer);
}

void gfx_imm_push_matrix(ImmRenderer* renderer, const mat4* matrix) {
  assert(renderer);
  assert(matrix);
  assert(renderer->stack_pointer >= 0);
  ASSERT(renderer->stack_pointer < IMM_MAX_NUM_MATRICES);

  renderer->stack_pointer += 1;
  renderer->matrix_stack[renderer->stack_pointer] =
      mat4_mul(*matrix, renderer->matrix_stack[renderer->stack_pointer - 1]);

  update_shader_model_matrix(renderer);
}

void gfx_imm_pop_matrix(ImmRenderer* renderer) {
  assert(renderer);
  ASSERT(renderer->stack_pointer > 0);

  // For debugging, zero out the matrix stack as matrices are popped out.
  memset(
      &renderer->matrix_stack[renderer->stack_pointer], 0,
      sizeof(renderer->matrix_stack[0]));
  renderer->stack_pointer -= 1;

  update_shader_model_matrix(renderer);
}

void gfx_imm_translate(ImmRenderer* renderer, vec3 offset) {
  assert(renderer);

  const mat4 mat = mat4_translate(offset);
  gfx_imm_push_matrix(renderer, &mat);
}

void gfx_imm_set_model_matrix(ImmRenderer* renderer, const mat4* model) {
  assert(renderer);
  assert(model);

  gfx_imm_flush(renderer);

  renderer->matrix_stack[0] = *model;
  renderer->stack_pointer   = 0;
  update_shader_model_matrix(renderer);
}

void gfx_imm_set_camera(ImmRenderer* renderer, const Camera* camera) {
  assert(renderer);
  assert(renderer->shader);

  gfx_imm_flush(renderer);

  const mat4 view = spatial3_inverse_transform(&camera->spatial);
  // const mat4 view_proj = mat4_mul(camera->projection, view);
  // gfx_imm_set_view_projection_matrix(renderer, &view_proj);
  renderer->view           = view;
  renderer->projection     = camera->projection;
  renderer->camera_changed = true;
}

// void gfx_imm_set_view_projection_matrix(
//     ImmRenderer* renderer, const mat4* view_proj) {
//   assert(renderer);
//   assert(renderer->shader);
//
//   gfx_imm_flush(renderer);
//   gfx_set_mat4_uniform(renderer->shader, "ViewProjection", view_proj);
// }