#include "renderer_impl.h"

#include "scene/animation_impl.h"
#include "scene/camera_impl.h"
#include "scene/light_impl.h"
#include "scene/material_impl.h"
#include "scene/mesh_impl.h"
#include "scene/model_impl.h"
#include "scene/node_impl.h"
#include "scene/object_impl.h"
#include "scene/scene_impl.h"
#include "scene/scene_memory.h"

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

#include <log/log.h>
#include <math/mat4.h>
#include <math/spatial3.h>

#include <assert.h>

// TODO: Move to a header like "constants.h".
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;

bool renderer_make(Renderer* renderer, GfxCore* gfxcore) {
  assert(renderer);
  assert(gfxcore);

  renderer->gfxcore = gfxcore;

  return true;
}

void renderer_destroy(Renderer* renderer) {
  if (!renderer) {
    return;
  }
  assert(renderer->gfxcore);
  GfxCore* gfxcore = renderer->gfxcore;
  if (renderer->ibl) {
    gfx_destroy_ibl(gfxcore, &renderer->ibl);
  }
  if (renderer->shaders.debug) {
    gfx_destroy_shader_program(gfxcore, &renderer->shaders.debug);
  }
  if (renderer->shaders.normals) {
    gfx_destroy_shader_program(gfxcore, &renderer->shaders.normals);
  }
  if (renderer->shaders.normal_mapped_normals) {
    gfx_destroy_shader_program(
        gfxcore, &renderer->shaders.normal_mapped_normals);
  }
  if (renderer->shaders.tangents) {
    gfx_destroy_shader_program(gfxcore, &renderer->shaders.tangents);
  }
}

/// Initialize renderer state for IBL if not already initialized.
static bool init_ibl(Renderer* renderer) {
  assert(renderer);

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

  if (!renderer->brdf_integration_map &&
      !(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;
}

static ShaderProgram* load_shader(Renderer* renderer, RenderSceneMode mode) {
  assert(renderer);

#define LOAD_AND_RETURN(pShader, constructor)   \
  {                                             \
    if (!pShader) {                             \
      pShader = constructor(renderer->gfxcore); \
    }                                           \
    assert(pShader);                            \
    return pShader;                             \
  }

  switch (mode) {
  case RenderDefault:
    return 0;
  case RenderDebug:
    LOAD_AND_RETURN(renderer->shaders.debug, gfx_make_debug3d_shader);
  case RenderNormals:
    LOAD_AND_RETURN(renderer->shaders.normals, gfx_make_view_normals_shader);
  case RenderNormalMappedNormals:
    LOAD_AND_RETURN(
        renderer->shaders.normal_mapped_normals,
        gfx_make_view_normal_mapped_normals_shader);
  case RenderTangents:
    LOAD_AND_RETURN(renderer->shaders.tangents, gfx_make_view_tangents_shader);
  }
  assert(false);
  return 0;
}

// static void log_matrix(const mat4* m) {
//   for (int row = 0; row < 4; ++row) {
//     LOGI("[ %5.2f, %5.2f, %5.2f, %5.2f ]", m->val[0][row], m->val[1][row],
//          m->val[2][row], m->val[3][row]);
//   }
// }

/// Computes irradiance and prefiltered environment maps for the light if they
/// have not been already computed.
static bool setup_environment_light(
    Renderer* renderer, GfxCore* gfxcore, EnvironmentLight* light) {
  assert(renderer);
  assert(light);

  if (!init_ibl(renderer)) {
    return false;
  }

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

  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;
}

typedef struct RenderState {
  GfxCore*       gfxcore;
  Renderer*      renderer;
  ShaderProgram* shader; // Null to use scene shaders.
  const Scene*   scene;
  const Camera*  camera;
  const mat4*    camera_rotation; // From camera to world space, rotation only.
  const mat4*    view_matrix;
  const mat4*    projection;
  const float    fovy;
  const float    aspect;
  Light*         environment_light;
  const Anima*   anima;
  size_t         num_joints;
  mat4           joint_matrices[GFX_MAX_NUM_JOINTS];
} RenderState;

/// Load joint matrices into the render state.
static void load_skeleton(RenderState* state, skeleton_idx skeleton_index) {
  assert(state);
  assert(skeleton_index.val != 0);

  const Skeleton* skeleton = mem_get_skeleton(skeleton_index);
  assert(skeleton);
  assert(skeleton->num_joints <= GFX_MAX_NUM_JOINTS);

  state->num_joints = skeleton->num_joints;

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

/// Draw the scene recursively.
static void draw_recursively(
    RenderState* state, mat4 parent_transform, const SceneNode* node) {
  assert(state);
  const mat4 node_transform = mat4_mul(parent_transform, node->transform);

  // Anima.
  if (node->type == AnimaNode) {
    state->anima = gfx_get_node_anima(node);
  }
  // Activate light.
  else if (node->type == LightNode) {
    Light* light = mem_get_light(node->light);
    assert(light);

    if (light->type == EnvironmentLightType) {
      bool result = setup_environment_light(
          state->renderer, state->gfxcore, &light->environment);
      // TODO: Handle the result in a better way.
      assert(result);
      state->environment_light = light;
    }
  }
  // Model.
  else if (node->type == ModelNode) {
    const Model*     model = gfx_get_node_model(node);
    const SceneNode* root  = mem_get_node(model->root);
    draw_recursively(state, parent_transform, root);
  }
  // Render object.
  else if (node->type == ObjectNode) {
    const SceneObject* object = mem_get_object(node->object);
    assert(object);

    // TODO: Here we would frustum-cull the object.

    // TODO: Avoid computing matrices like Modelview or MVP if the shader does
    // not use them.
    const mat4 model_matrix = node_transform;
    const mat4 modelview    = mat4_mul(*state->view_matrix, model_matrix);
    const mat4 mvp          = mat4_mul(*state->projection, modelview);

    if (object->skeleton.val) {
      load_skeleton(state, object->skeleton);
    }

    for (mesh_link_idx mesh_link_index = object->mesh_link;
         mesh_link_index.val;) {
      const MeshLink* mesh_link = mem_get_mesh_link(mesh_link_index);
      mesh_link_index           = mesh_link->next;

      const Mesh* mesh = mem_get_mesh(mesh_link->mesh);
      if (!mesh) {
        continue;
      }
      assert(mesh->geometry);
      assert(mesh->material);

      // TODO: Here we would frustum-cull the mesh. The AABB would have to be
      // transformed by the model matrix. Rotation would make the AABB
      // relatively large, but still, the culling would be conservative.

      // Apply common shader uniforms not captured by materials.
      ShaderProgram* shader = state->shader ? state->shader : mesh->shader;
      gfx_set_mat4_uniform(shader, "ModelMatrix", &model_matrix);
      gfx_set_mat4_uniform(shader, "Modelview", &modelview);
      gfx_set_mat4_uniform(shader, "View", state->view_matrix);
      gfx_set_mat4_uniform(shader, "Projection", state->projection);
      gfx_set_mat4_uniform(shader, "MVP", &mvp);
      gfx_set_mat4_uniform(shader, "CameraRotation", state->camera_rotation);
      gfx_set_float_uniform(shader, "Fovy", state->fovy);
      gfx_set_float_uniform(shader, "Aspect", state->aspect);
      if (state->camera) {
        gfx_set_vec3_uniform(
            shader, "CameraPosition", state->camera->spatial.p);
      }
      if (state->num_joints > 0) {
        gfx_set_mat4_array_uniform(
            shader, "JointMatrices", state->joint_matrices, state->num_joints);
      }
      // Apply lights.
      if (state->environment_light) {
        const EnvironmentLight* light = &state->environment_light->environment;
        assert(light->environment_map);
        assert(light->irradiance_map);
        assert(light->prefiltered_environment_map);
        assert(state->renderer->brdf_integration_map);
        gfx_set_texture_uniform(
            shader, "BRDFIntegrationMap",
            state->renderer->brdf_integration_map);
        gfx_set_texture_uniform(shader, "Sky", light->environment_map);
        gfx_set_texture_uniform(shader, "IrradianceMap", light->irradiance_map);
        gfx_set_texture_uniform(
            shader, "PrefilteredEnvironmentMap",
            light->prefiltered_environment_map);
        gfx_set_float_uniform(
            shader, "MaxReflectionLOD", light->max_reflection_lod);
      }
      material_activate(shader, mesh->material);
      gfx_activate_shader_program(shader);
      gfx_apply_uniforms(shader);
      gfx_render_geometry(mesh->geometry);
    }

    // Reset state for next object.
    state->num_joints = 0;
  }

  // Render children recursively.
  for (node_idx child_index = node->child; child_index.val;) {
    const SceneNode* child = mem_get_node(child_index);
    draw_recursively(state, node_transform, child);
    child_index = child->next;
  }
}

void gfx_render_scene(Renderer* renderer, const RenderSceneParams* params) {
  assert(renderer);
  assert(params);
  assert(params->scene);

  ShaderProgram* const shader = load_shader(renderer, params->mode);

  const Scene*       scene  = params->scene;
  const SceneCamera* camera = params->camera;

  GfxCore* gfxcore = renderer->gfxcore;

  mat4 projection, camera_rotation, view_matrix;
  if (camera) {
    projection = camera->camera.projection;
    camera_rotation =
        mat4_rotation(spatial3_transform(&camera->camera.spatial));
    view_matrix = spatial3_inverse_transform(&camera->camera.spatial);
  } else {
    projection      = mat4_id();
    camera_rotation = mat4_id();
    view_matrix     = mat4_id();
  }

  int x, y, width, height;
  gfx_get_viewport(gfxcore, &x, &y, &width, &height);
  const float aspect = (float)width / (float)height;

  RenderState state = {
      .gfxcore           = gfxcore,
      .renderer          = renderer,
      .shader            = shader,
      .scene             = scene,
      .camera            = &camera->camera,
      .camera_rotation   = &camera_rotation,
      .view_matrix       = &view_matrix,
      .projection        = &projection,
      .environment_light = 0,
      // Assuming a perspective matrix.
      .fovy   = atan(1.0 / (mat4_at(projection, 1, 1))) * 2,
      .aspect = aspect};

  draw_recursively(&state, mat4_id(), scene->root);
}

static void update_rec(SceneNode* node, const SceneCamera* camera, R t) {
  assert(node);
  assert(camera);

  const NodeType node_type = gfx_get_node_type(node);

  // TODO: Models do not need to be animated if they are not visible to the
  //  camera.
  if (node_type == AnimaNode) {
    Anima* anima = gfx_get_node_anima_mut(node);
    gfx_update_animation(anima, (R)t);
  } else if (node_type == ModelNode) {
    Model*     model = gfx_get_node_model_mut(node);
    SceneNode* root  = gfx_get_model_root_mut(model);
    update_rec(root, camera, t);
  }

  // Children.
  SceneNode* child = gfx_get_node_child_mut(node);
  while (child) {
    update_rec(child, camera, t);
    child = gfx_get_node_sibling_mut(child);
  }
}

void gfx_update(Scene* scene, const SceneCamera* camera, R t) {
  assert(scene);
  assert(camera);

  SceneNode* node = gfx_get_scene_root(scene);
  update_rec(node, camera, t);
}