#include <gfx/util/ibl.h>

#include <gfx/core.h>
#include <gfx/util/geometry.h>
#include <gfx/util/shader.h>
#include <math/mat4.h>

#include <assert.h>
#include <stdlib.h>

typedef struct IBL {
  Geometry*      quad;
  ShaderProgram* brdf_integration_map_shader;
  ShaderProgram* irradiance_map_shader;
  ShaderProgram* prefiltered_environment_map_shader;
  Texture*       brdf_integration_map;
  FrameBuffer*   framebuffer;
  mat4           rotations[6];
} IBL;

static const CubemapFace faces[6] = {
    CubemapFacePosX, // Right.
    CubemapFaceNegX, // Left.
    CubemapFacePosY, // Up.
    CubemapFaceNegY, // Down.
    CubemapFacePosZ, // Back.
    CubemapFaceNegZ, // Front.
};

static const float flips[6] = {
    -1.0f, // Right.
    -1.0f, // Left.
    +1.0f, // Up.
    +1.0f, // Down.
    -1.0f, // Back.
    -1.0f, // Front.
};

IBL* gfx_make_ibl(GfxCore* gfxcore) {
  assert(gfxcore);

  IBL* ibl = calloc(1, sizeof(IBL));
  if (!ibl) {
    return 0;
  }

  if (!(ibl->quad = gfx_make_quad_11(gfxcore))) {
    goto cleanup;
  }

  // We only need the BRDF integration once since we are caching the map, but
  // compiling the shader up front may lead to fewer surprises. Not that the
  // shader is fully compiled up front anyway, since the driver will typically
  // defer full compilation to the first draw call.
  if (!(ibl->brdf_integration_map_shader =
            gfx_make_brdf_integration_map_shader(gfxcore))) {
    goto cleanup;
  }

  if (!(ibl->irradiance_map_shader = gfx_make_irradiance_map_shader(gfxcore))) {
    goto cleanup;
  }

  if (!(ibl->prefiltered_environment_map_shader =
            gfx_make_prefiltered_environment_map_shader(gfxcore))) {
    goto cleanup;
  }

  // Create an empty framebuffer for now. Will attach the colour buffer later
  // as we render the faces of the cube.
  if (!(ibl->framebuffer = gfx_make_framebuffer(
            gfxcore,
            &(FrameBufferDesc){
                .colour =
                    (FrameBufferAttachment){.type = FrameBufferNoAttachment},
                .depth = (FrameBufferAttachment){
                    .type = FrameBufferNoAttachment}}))) {
    goto cleanup;
  }

  // TODO: Debug the camera rotations. Irradiance debug output should appear
  // just like the input cubemap.

  // Right.
  ibl->rotations[0] = mat4_lookat(
      /*position=*/vec3_make(0, 0, 0),
      /*target=*/vec3_make(1, 0, 0),
      /*up=*/vec3_make(0, 1, 0));
  // Left.
  ibl->rotations[1] = mat4_lookat(
      /*position=*/vec3_make(0, 0, 0),
      /*target=*/vec3_make(-1, 0, 0),
      /*up=*/vec3_make(0, 1, 0));
  // Up.
  ibl->rotations[2] = mat4_lookat(
      /*position=*/vec3_make(0, 0, 0),
      /*target=*/vec3_make(0, 1, 0),
      /*up=*/vec3_make(0, 0, 1));
  // Down.
  ibl->rotations[3] = mat4_lookat(
      /*position=*/vec3_make(0, 0, 0),
      /*target=*/vec3_make(0, -1, 0),
      /*up=*/vec3_make(0, 0, -1));
  // Back.
  ibl->rotations[4] = mat4_lookat(
      /*position=*/vec3_make(0, 0, 0),
      /*target=*/vec3_make(0, 0, 1),
      /*up=*/vec3_make(0, 1, 0));
  // Front.
  ibl->rotations[5] = mat4_lookat(
      /*position=*/vec3_make(0, 0, 0),
      /*target=*/vec3_make(0, 0, -1),
      /*up=*/vec3_make(0, 1, 0));

  return ibl;

cleanup:
  gfx_destroy_ibl(gfxcore, &ibl);
  return 0;
}

void gfx_destroy_ibl(GfxCore* gfxcore, IBL** ibl) {
  if (!ibl) {
    return;
  }
  if ((*ibl)->quad) {
    gfx_destroy_geometry(gfxcore, &(*ibl)->quad);
  }
  if ((*ibl)->brdf_integration_map_shader) {
    gfx_destroy_shader_program(gfxcore, &(*ibl)->brdf_integration_map_shader);
  }
  if ((*ibl)->irradiance_map_shader) {
    gfx_destroy_shader_program(gfxcore, &(*ibl)->irradiance_map_shader);
  }
  if ((*ibl)->prefiltered_environment_map_shader) {
    gfx_destroy_shader_program(
        gfxcore, &(*ibl)->prefiltered_environment_map_shader);
  }
  if ((*ibl)->brdf_integration_map) {
    gfx_destroy_texture(gfxcore, &(*ibl)->brdf_integration_map);
  }
  if ((*ibl)->framebuffer) {
    gfx_destroy_framebuffer(gfxcore, &(*ibl)->framebuffer);
  }
  free(*ibl);
  *ibl = 0;
}

Texture* gfx_make_brdf_integration_map(
    IBL* ibl, GfxCore* gfxcore, int width, int height) {
  assert(ibl);
  assert(gfxcore);

  if (ibl->brdf_integration_map) {
    return ibl->brdf_integration_map;
  }

  bool success = false;

  if (!(ibl->brdf_integration_map = gfx_make_texture(
            gfxcore, &(TextureDesc){
                         .width     = width,
                         .height    = height,
                         .depth     = 1,
                         .dimension = Texture2D,
                         .format    = TextureRG16F,
                         .filtering = LinearFiltering,
                         .wrap      = ClampToEdge,
                         .mipmaps   = false}))) {
    goto cleanup;
  }

  gfx_activate_framebuffer(ibl->framebuffer);
  gfx_framebuffer_set_viewport(ibl->framebuffer, 0, 0, width, height);
  gfx_activate_shader_program(ibl->brdf_integration_map_shader);
  if (!gfx_framebuffer_attach_colour(
          ibl->framebuffer, &(FrameBufferAttachment){
                                .type              = FrameBufferTexture,
                                .texture.texture   = ibl->brdf_integration_map,
                                .texture.mip_level = 0})) {
    goto cleanup;
  }
  gfx_render_geometry(ibl->quad);

  success = true;

cleanup:
  gfx_deactivate_shader_program(ibl->brdf_integration_map_shader);
  gfx_deactivate_framebuffer(ibl->framebuffer);
  if (!success && ibl->brdf_integration_map) {
    gfx_destroy_texture(gfxcore, &ibl->brdf_integration_map);
    return 0;
  } else {
    return ibl->brdf_integration_map;
  }
}

Texture* gfx_make_irradiance_map(
    IBL* ibl, GfxCore* gfxcore, const Texture* environment_map, int width,
    int height) {
  assert(ibl);
  assert(gfxcore);
  assert(environment_map);

  bool success = false;

  Texture* irradiance_map = 0;

  // TODO: Could define colour-renderable texture formats separately to make
  // framebuffer creation less error-prone. Or, at the very least, validate the
  // choice at runtime.
  //
  // Make sure to use a float colour format to avoid [0,1] clamping when the
  // irradiance values are computed!
  if (!(irradiance_map = gfx_make_texture(
            gfxcore, &(TextureDesc){
                         .width     = width,
                         .height    = height,
                         .depth     = 1,
                         .dimension = TextureCubeMap,
                         .format    = TextureR11G11B10F,
                         .filtering = LinearFiltering,
                         .mipmaps   = false}))) {
    goto cleanup;
  }

  gfx_activate_framebuffer(ibl->framebuffer);
  gfx_framebuffer_set_viewport(ibl->framebuffer, 0, 0, width, height);
  gfx_activate_shader_program(ibl->irradiance_map_shader);
  gfx_set_texture_uniform(ibl->irradiance_map_shader, "Sky", environment_map);
  for (int i = 0; i < 6; ++i) {
    if (!gfx_framebuffer_attach_colour(
            ibl->framebuffer, &(FrameBufferAttachment){
                                  .type            = FrameBufferCubemapTexture,
                                  .cubemap.face    = faces[i],
                                  .cubemap.texture = irradiance_map})) {
      goto cleanup;
    }
    gfx_set_float_uniform(ibl->irradiance_map_shader, "Flip", flips[i]);
    gfx_set_mat4_uniform(
        ibl->irradiance_map_shader, "CameraRotation", &ibl->rotations[i]);
    gfx_apply_uniforms(ibl->irradiance_map_shader);
    gfx_render_geometry(ibl->quad);
  }

  success = true;

cleanup:
  gfx_deactivate_shader_program(ibl->irradiance_map_shader);
  gfx_deactivate_framebuffer(ibl->framebuffer);
  if (!success && irradiance_map) {
    gfx_destroy_texture(gfxcore, &irradiance_map);
    return 0;
  } else {
    return irradiance_map;
  }
}

Texture* gfx_make_prefiltered_environment_map(
    IBL* ibl, GfxCore* gfxcore, const Texture* environment_map, int width,
    int height, int* max_mip_level) {
  assert(ibl);
  assert(gfxcore);
  assert(environment_map);
  assert(max_mip_level);

  bool success = false;

  Texture* prefiltered_env_map = 0;

  if (!(prefiltered_env_map = gfx_make_texture(
            gfxcore, &(TextureDesc){
                         .width     = width,
                         .height    = height,
                         .depth     = 1,
                         .dimension = TextureCubeMap,
                         .format    = TextureR11G11B10F,
                         .filtering = LinearFiltering,
                         .mipmaps   = true}))) {
    goto cleanup;
  }

  gfx_activate_framebuffer(ibl->framebuffer);
  gfx_activate_shader_program(ibl->prefiltered_environment_map_shader);
  gfx_set_texture_uniform(
      ibl->prefiltered_environment_map_shader, "Sky", environment_map);
  const int max_mip = (int)(rlog2(min(width, height)));
  for (int mip = 0; mip <= max_mip; ++mip) {
    const int   mip_width  = width >> mip;
    const int   mip_height = height >> mip;
    const float roughness  = (float)mip / (float)(max_mip);
    gfx_framebuffer_set_viewport(ibl->framebuffer, 0, 0, mip_width, mip_height);
    gfx_set_float_uniform(
        ibl->prefiltered_environment_map_shader, "Roughness", roughness);

    for (int i = 0; i < 6; ++i) {
      if (!gfx_framebuffer_attach_colour(
              ibl->framebuffer, &(FrameBufferAttachment){
                                    .type         = FrameBufferCubemapTexture,
                                    .cubemap.face = faces[i],
                                    .cubemap.mip_level = mip,
                                    .cubemap.texture = prefiltered_env_map})) {
        goto cleanup;
      }
      gfx_set_float_uniform(
          ibl->prefiltered_environment_map_shader, "Flip", flips[i]);
      gfx_set_mat4_uniform(
          ibl->prefiltered_environment_map_shader, "CameraRotation",
          &ibl->rotations[i]);
      gfx_apply_uniforms(ibl->prefiltered_environment_map_shader);
      gfx_render_geometry(ibl->quad);
    }
  }

  *max_mip_level = max_mip;

  success = true;

cleanup:
  gfx_deactivate_shader_program(ibl->prefiltered_environment_map_shader);
  gfx_deactivate_framebuffer(ibl->framebuffer);
  if (!success && prefiltered_env_map) {
    gfx_destroy_texture(gfxcore, &prefiltered_env_map);
    return 0;
  } else {
    return prefiltered_env_map;
  }
}