/*
Copyright 2022 Marc Sunet

Redistribution and use in source and binary forms, with or without modification,
are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice, this
list of conditions and the following disclaimer.

2. Redistributions in binary form must reproduce the above copyright notice,
this list of conditions and the following disclaimer in the documentation and/or
other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/
#include "cgltf_tangents.h"
#include "cgltf.h"

#include "MikkTSpace/mikktspace.h"

#include <assert.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#ifdef CGLTF_TANGENTS_DEBUG
#include <stdio.h>
#define DLOG printf
#else
#define DLOG(...)
#endif

#include <stdio.h> // TODO: Remove me.

#define CGLTF_OPTIONS_MALLOC(size)                                             \
  options->memory.alloc(options->memory.user_data, size)

#define CGLTF_OPTIONS_FREE(ptr)                                                \
  options->memory.free(options->memory.user_data, ptr)

static void* cgltf_default_alloc(void* user, cgltf_size size) {
  (void)user;
  return malloc(size);
}

static void cgltf_default_free(void* user, void* ptr) {
  (void)user;
  free(ptr);
}

static const cgltf_size NUM_TANGENT_COMPONENTS = 4; // X,Y,Z,fSign

static float normalize_i8(int8_t x) { return (float)x / 128.0; }
static float normalize_u8(uint8_t x) { return (float)x / 255.0; }
static float normalize_i16(int16_t x) { return (float)x / 32768.0; }
static float normalize_u16(uint16_t x) { return (float)x / 65535.0; }
static float normalize_u32(uint32_t x) { return (float)x / 4294967295.0; }

static cgltf_size num_vertex_attrib_components(cgltf_type type) {
  switch (type) {
  case cgltf_type_scalar:
    return 1;
  case cgltf_type_vec2:
    return 2;
  case cgltf_type_vec3:
    return 3;
  case cgltf_type_vec4:
    return 4;
  default:
    assert(false);
    return 0;
  }
}

static cgltf_size cgltf_component_type_size_bytes(cgltf_component_type type) {
  switch (type) {
  case cgltf_component_type_r_8:
    return 1;
  case cgltf_component_type_r_8u:
    return 1;
  case cgltf_component_type_r_16:
    return 2;
  case cgltf_component_type_r_16u:
    return 2;
  case cgltf_component_type_r_32u:
    return 4;
  case cgltf_component_type_r_32f:
    return 4;
  default:
    assert(false);
    return 0;
  }
}

static cgltf_size default_stride(cgltf_type type,
                                 cgltf_component_type component_type) {
  return num_vertex_attrib_components(type) *
         cgltf_component_type_size_bytes(component_type);
}

// -----------------------------------------------------------------------------
// MikkTSpace interface

// An array of values for a given vertex attribute or for vertex indices.
// For positions and normals, glTF mandates floats.
// Texcoords and indices can be different types and vary in size: 8-bit, 16-bit,
// or 32-bit.
// We store void* pointers so that we can do byte pointer arithmetic.
typedef struct Buffer {
  const void* start;         // X-coordinate of the first attribute.
  const void* end;           // One byte past the end of the buffer.
  cgltf_size stride_bytes;   // Stride in bytes between each value.
  cgltf_component_type type; // Type of each value in the buffer.
} Buffer;

// User data for mesh processing.
// See: https://registry.khronos.org/glTF/specs/2.0/glTF-2.0.html#meshes
// Buffer pointers have the accessor + view offsets baked in so that we do the
// addition only once.
typedef struct SMikkUserData {
  const cgltf_primitive* primitive;
  // Index buffer may be empty (mesh primitive has no indices).
  Buffer indices;
  // Vertex attributes.
  Buffer positions;
  Buffer normals;
  Buffer texcoords;
  // Output tangents.
  void* tangents;
} SMikkUserData;

static cgltf_size get_vertex_index(const SMikkUserData* data, cgltf_size iFace,
                                   cgltf_size iVert) {
  const cgltf_primitive* primitive = data->primitive;

  // First compute a vertex index as if the mesh primitive had no indices.
  cgltf_size vertex_idx = 0;
  switch (primitive->type) {
  case cgltf_primitive_type_triangles:
    vertex_idx = iFace * 3 + iVert;
    break;
  case cgltf_primitive_type_triangle_strip:
    // For triangle strips:
    // face 0 -> verts 0, 1, 2
    // face 1 -> verts 1, 3, 2 (1, 2, 3 flipped)
    // face 2 -> verts 2, 3, 4
    // face 3 -> verts 3, 5, 4 (3, 4, 5 flipped)
    // ...
    // face N=2k   -> verts N, N+1, N+2
    // face N=2k+1 -> verts N, N+2, N+1
    if (iFace & 1) {
      // Flip the winding of odd faces so that the is consistent with the even
      // ones.
      // iVert = 0 -> vert 0
      // iVert = 1 -> vert 2
      // iVert = 2 -> vert 1
      vertex_idx = iFace + (2 - iVert);
    } else {
      vertex_idx = iFace + iVert;
    }
    break;
  case cgltf_primitive_type_triangle_fan:
    // For triangle fans:
    // face 0 -> verts 0, 1, 2
    // face 1 -> verts 0, 2, 3
    // face 2 -> verts 0, 3, 4
    // face 3 -> verts 0, 4, 5
    // ...
    // face N -> verts 0, N=1, N=2
    if (iVert == 0) {
      vertex_idx = 0;
    } else {
      vertex_idx = iFace + iVert;
    }
    break;
  default:
    assert(false);
    break;
  }

  // If the mesh primitive has vertex indices, then vertex_idx is actually the
  // index of the index. Index the index buffer with vertex_idx to find the
  // real vertex index.
  if (primitive->indices != NULL) {
    const void* p_idx =
        data->indices.start + vertex_idx * data->indices.stride_bytes;
    switch (data->indices.type) {
    case cgltf_component_type_r_8:
      vertex_idx = *((int8_t*)p_idx);
      break;
    case cgltf_component_type_r_8u:
      vertex_idx = *((uint8_t*)p_idx);
      break;
    case cgltf_component_type_r_16:
      vertex_idx = *((int16_t*)p_idx);
      break;
    case cgltf_component_type_r_16u:
      vertex_idx = *((uint16_t*)p_idx);
      break;
    case cgltf_component_type_r_32u:
      vertex_idx = *((uint32_t*)p_idx);
      break;
    default:
      assert(false);
      break;
    }
  }

  return vertex_idx;
}

static const void* get_vertex(const Buffer* buffer, cgltf_size index) {
  // Stride is the offset in bytes between vertex attributes.
  const void* vertex = buffer->start + buffer->stride_bytes * index;
  assert(vertex < buffer->end);
  return vertex;
}

static const void* get_position(const SMikkUserData* data, cgltf_size index) {
  return get_vertex(&data->positions, index);
}

static const void* get_normal(const SMikkUserData* data, cgltf_size index) {
  return get_vertex(&data->normals, index);
}

static const void* get_texcoord(const SMikkUserData* data, cgltf_size index) {
  return get_vertex(&data->texcoords, index);
}

static float* get_tangent(void* buffer, cgltf_size index) {
  // Tangents are tightly packed.
  return (float*)(buffer) + NUM_TANGENT_COMPONENTS * index;
}

static int SMikk_get_num_faces(const SMikkTSpaceContext* pContext) {
  SMikkUserData* data = (SMikkUserData*)pContext->m_pUserData;
  const cgltf_primitive* primitive = data->primitive;

  // Find the number of effective vertices (vertices or indices) in the mesh
  // primitive.
  //
  // https://registry.khronos.org/glTF/specs/2.0/glTF-2.0.html#meshes
  //
  // "All attribute accessors for a given primitive MUST have the same count.
  // When indices property is not defined, attribute accessors' count indicates
  // the number of vertices to render; when indices property is defined, it
  // indicates the upper (exclusive) bound on the index values in the indices
  // accessor, i.e., all index values MUST be less than attribute accessors'
  // count."
  const cgltf_size num_verts = (primitive->indices != NULL)
                                   ? primitive->indices->count
                                   : primitive->attributes_count;

  // Determine the number of faces given the number of vertices.
  //
  // We use the fact that glTF only supports triangles for faces.
  // https://registry.khronos.org/glTF/specs/2.0/glTF-2.0.html#meshes
  switch (primitive->type) {
  case cgltf_primitive_type_triangles:
    return (int)num_verts / 3;
  case cgltf_primitive_type_triangle_strip:
  case cgltf_primitive_type_triangle_fan:
    return (int)num_verts - 2;
  default:
    return 0;
  }
}

int SMikk_get_num_vertices_of_face(const SMikkTSpaceContext* pContext,
                                   const int iFace) {
  // Triangles are the only faces supported by glTF.
  // https://registry.khronos.org/glTF/specs/2.0/glTF-2.0.html#meshes
  return 3;
}

void SMikk_get_position(const SMikkTSpaceContext* pContext, float fvPosOut[],
                        const int iFace, const int iVert) {
  const SMikkUserData* data = (SMikkUserData*)pContext->m_pUserData;
  const cgltf_primitive* primitive = data->primitive;

  const cgltf_size idx = get_vertex_index(data, iFace, iVert);
  const float* coord = get_position(data, idx);
  fvPosOut[0] = *coord++;
  fvPosOut[1] = *coord++;
  fvPosOut[2] = *coord;
  DLOG("Position (face: %d, vert: %d): %f, %f, %f; idx: %lu\n", iFace, iVert,
       fvPosOut[0], fvPosOut[1], fvPosOut[2], idx);
}

void SMikk_get_normal(const SMikkTSpaceContext* pContext, float fvNormOut[],
                      const int iFace, const int iVert) {
  const SMikkUserData* data = (SMikkUserData*)pContext->m_pUserData;
  const cgltf_primitive* primitive = data->primitive;

  const cgltf_size idx = get_vertex_index(data, iFace, iVert);
  const float* coord = get_normal(data, idx);
  fvNormOut[0] = *coord++;
  fvNormOut[1] = *coord++;
  fvNormOut[2] = *coord;
  DLOG("Normal (face: %d, vert: %d): %f, %f, %f\n", iFace, iVert, fvNormOut[0],
       fvNormOut[1], fvNormOut[2]);
}

void SMikk_get_texcoord(const SMikkTSpaceContext* pContext, float fvTexcOut[],
                        const int iFace, const int iVert) {
  const SMikkUserData* data = (SMikkUserData*)pContext->m_pUserData;
  const cgltf_primitive* primitive = data->primitive;

  const cgltf_size idx = get_vertex_index(data, iFace, iVert);
  const void* coord = get_texcoord(data, idx);
  switch (data->texcoords.type) {
  case cgltf_component_type_r_8: {
    const int8_t* c = coord;
    fvTexcOut[0] = normalize_i8(*c++);
    fvTexcOut[1] = normalize_i8(*c);
    break;
  }
  case cgltf_component_type_r_8u: {
    const uint8_t* c = coord;
    fvTexcOut[0] = normalize_u8(*c++);
    fvTexcOut[1] = normalize_u8(*c);
    break;
  }
  case cgltf_component_type_r_16: {
    const int16_t* c = coord;
    fvTexcOut[0] = normalize_i16(*c++);
    fvTexcOut[1] = normalize_i16(*c);
    break;
  }
  case cgltf_component_type_r_16u: {
    const uint16_t* c = coord;
    fvTexcOut[0] = normalize_u16(*c++);
    fvTexcOut[1] = normalize_u16(*c);
    break;
  }
  case cgltf_component_type_r_32u: {
    const uint32_t* c = coord;
    fvTexcOut[0] = normalize_u32(*c++);
    fvTexcOut[1] = normalize_u32(*c);
    break;
  }
  case cgltf_component_type_r_32f: {
    const float* c = coord;
    fvTexcOut[0] = *c++;
    fvTexcOut[1] = *c;
    break;
  }
  default:
    assert(false);
    break;
  }
  DLOG("Texcoord (face: %d, vert: %d): %f, %f\n", iFace, iVert, fvTexcOut[0],
       fvTexcOut[1]);
}

void SMikk_set_TSpace_basic(const SMikkTSpaceContext* pContext,
                            const float fvTangent[], const float fSign,
                            const int iFace, const int iVert) {
  SMikkUserData* data = (SMikkUserData*)pContext->m_pUserData;
  const cgltf_primitive* primitive = data->primitive;

  const cgltf_size idx = get_vertex_index(data, iFace, iVert);
  float* coord = get_tangent(data->tangents, idx);
  *coord++ = fvTangent[0];
  *coord++ = fvTangent[1];
  *coord++ = fvTangent[2];
  *coord = fSign;
  DLOG("Tangent (face: %d, vert: %d): %f, %f, %f; sign: %f\n", iFace, iVert,
       fvTangent[0], fvTangent[1], fvTangent[2], fSign);
}

// -----------------------------------------------------------------------------

static bool has_normal_map(const cgltf_primitive* primitive) {
  return (primitive->material != NULL) &&
         (primitive->material->normal_texture.texture != NULL);
}

static const cgltf_attribute* find_attribute(const cgltf_primitive* primitive,
                                             cgltf_attribute_type type) {
  for (cgltf_size i = 0; i < primitive->attributes_count; ++i) {
    const cgltf_attribute* attrib = &primitive->attributes[i];
    if (attrib->type == type) {
      return attrib;
    }
  }
  return NULL;
}

static bool has_attribute(const cgltf_primitive* primitive,
                          cgltf_attribute_type type) {
  return find_attribute(primitive, type) != NULL;
}

static bool has_positions3d(const cgltf_primitive* primitive) {
  const cgltf_attribute* attrib =
      find_attribute(primitive, cgltf_attribute_type_position);
  if (attrib) {
    return attrib->data->type == cgltf_type_vec3;
  }
  return false;
}

static bool has_normals(const cgltf_primitive* primitive) {
  return has_attribute(primitive, cgltf_attribute_type_normal);
}

static bool has_texcoords(const cgltf_primitive* primitive) {
  return has_attribute(primitive, cgltf_attribute_type_texcoord);
}

static bool has_tangents(const cgltf_primitive* primitive) {
  return has_attribute(primitive, cgltf_attribute_type_tangent);
}

static bool has_indices(const cgltf_primitive* primitive) {
  return primitive->indices != 0;
}

static cgltfTangentBuffer compute_tangents(const cgltf_options* options,
                                           const cgltf_data* data,
                                           cgltf_primitive* primitive) {
  cgltfTangentBuffer buffer = {0};
  SMikkUserData user = {0};
  cgltf_size num_verts = 0;

  user.primitive = primitive;

  if (primitive->indices != NULL) {
    const cgltf_accessor* accessor = primitive->indices;
    const cgltf_buffer_view* view = accessor->buffer_view;
    const cgltf_size offset_bytes = accessor->offset + view->offset;
    const void* buffer_data = view->buffer->data + offset_bytes;
    const void* buffer_end = view->buffer->data + view->offset + view->size;

    user.indices.start = buffer_data;
    user.indices.end = buffer_end;
    // Indices are tightly packed, stride 0.
    user.indices.stride_bytes =
        default_stride(accessor->type, accessor->component_type);
    user.indices.type = accessor->component_type;
  }

  for (cgltf_size i = 0; i < primitive->attributes_count; ++i) {
    const cgltf_attribute* attrib = &primitive->attributes[i];

    if ((attrib->type == cgltf_attribute_type_position) ||
        (attrib->type == cgltf_attribute_type_normal) ||
        (attrib->type == cgltf_attribute_type_texcoord)) {
      const cgltf_accessor* accessor = attrib->data;
      const cgltf_buffer_view* view = accessor->buffer_view;
      const cgltf_buffer* buffer = view->buffer;
      const cgltf_size offset_bytes = accessor->offset + view->offset;
      const cgltf_size stride_bytes =
          view->stride > 0
              ? view->stride
              : default_stride(accessor->type, accessor->component_type);
      // const cgltf_size size_bytes = view->size;
      const void* buffer_data = view->buffer->data + offset_bytes;
      const void* buffer_end = view->buffer->data + view->offset + view->size;

      Buffer* attrib_buffer = 0;

      if (attrib->type == cgltf_attribute_type_position) {
        // glTF currently mandates vec3 for positions. Caller should ensure
        // this.
        assert(accessor->type == cgltf_type_vec3);
        num_verts = attrib->data->count;
        attrib_buffer = &user.positions;
      } else if (attrib->type == cgltf_attribute_type_normal) {
        attrib_buffer = &user.normals;
      } else if (attrib->type == cgltf_attribute_type_texcoord) {
        attrib_buffer = &user.texcoords;
      }

      attrib_buffer->start = buffer_data;
      attrib_buffer->end = buffer_end;
      attrib_buffer->stride_bytes = stride_bytes;
      attrib_buffer->type = accessor->component_type;
    }
  }

  assert(user.positions.start);
  assert(user.positions.end);
  assert(user.normals.start);
  assert(user.normals.end);
  assert(user.texcoords.start);
  assert(user.texcoords.end);
  assert(num_verts > 0);

  const cgltf_size tangents_size_bytes =
      num_verts * NUM_TANGENT_COMPONENTS * sizeof(float);

  user.tangents = CGLTF_OPTIONS_MALLOC(tangents_size_bytes);
  if (!user.tangents) {
    return buffer;
  }

  SMikkTSpaceInterface interface = (SMikkTSpaceInterface){
      .m_getNumFaces = SMikk_get_num_faces,
      .m_getNumVerticesOfFace = SMikk_get_num_vertices_of_face,
      .m_getPosition = SMikk_get_position,
      .m_getNormal = SMikk_get_normal,
      .m_getTexCoord = SMikk_get_texcoord,
      .m_setTSpaceBasic = SMikk_set_TSpace_basic,
  };
  const SMikkTSpaceContext context = (SMikkTSpaceContext){
      .m_pInterface = &interface,
      .m_pUserData = &user,
  };
  if (!genTangSpaceDefault(&context)) {
    return buffer;
  }

  buffer.data = user.tangents;
  buffer.size_bytes = tangents_size_bytes;
  buffer.primitive = primitive;

  return buffer;
}

static void process_primitive(const cgltf_options* options,
                              const cgltf_data* data,
                              cgltf_primitive* primitive,
                              cgltfTangentBuffer* tangent_buffers,
                              cgltf_size* num_tangent_buffers) {
  DLOG("Processing primitive\n");
  cgltf_size cur_buffer = 0;
  // TODO: MikkTSpace should not be used with models with vertex indices. One
  // workaround is to unindex the mesh, compute tangents, and then re-index it.
  if (((primitive->type == cgltf_primitive_type_triangle_fan) ||
       (primitive->type == cgltf_primitive_type_triangle_strip) ||
       (primitive->type == cgltf_primitive_type_triangles)) &&
      has_normal_map(primitive) && !has_tangents(primitive) &&
      has_positions3d(primitive) && has_normals(primitive) &&
      has_texcoords(primitive) && !has_indices(primitive)) {
    *num_tangent_buffers += 1;
    if (tangent_buffers) {
      DLOG("Model with normal map missing tangents detected\n");
      tangent_buffers[cur_buffer] = compute_tangents(options, data, primitive);
      if (tangent_buffers[cur_buffer].data) {
        DLOG("Tangents computed\n");
      }
      cur_buffer++;
    }
  }
}

cgltf_result cgltf_compute_tangents(const cgltf_options* input_options,
                                    const cgltf_data* data,
                                    cgltfTangentBuffer** tangent_buffers,
                                    cgltf_size* num_tangent_buffers) {
  if ((input_options == NULL) || (data == NULL)) {
    return cgltf_result_invalid_options;
  }

  DLOG("cgltf_compute_tangents\n");

  cgltf_options options = *input_options;
  if (options.memory.alloc == NULL) {
    options.memory.alloc = &cgltf_default_alloc;
  }
  if (options.memory.free == NULL) {
    options.memory.free = &cgltf_default_free;
  }

  // First pass: compute the number of tangent buffers to be created.
  *num_tangent_buffers = 0;
  for (cgltf_size mesh_idx = 0; mesh_idx < data->meshes_count; ++mesh_idx) {
    const cgltf_mesh* mesh = &data->meshes[mesh_idx];

    for (cgltf_size prim_idx = 0; prim_idx < mesh->primitives_count;
         ++prim_idx) {
      // Pass in null for the tangent buffers to just compute the number of
      // buffers.
      process_primitive(&options, data, &mesh->primitives[prim_idx], 0,
                        num_tangent_buffers);
    }
  }
  DLOG("Number of primitives to be patched: %lu\n", *num_tangent_buffers);

  // Second pass: compute the tangents.
  if (*num_tangent_buffers > 0) {
    *tangent_buffers =
        options.memory.alloc(options.memory.user_data,
                             *num_tangent_buffers * sizeof(cgltfTangentBuffer));
    if (!*tangent_buffers) {
      return cgltf_result_out_of_memory;
    }

    cgltf_size tangent_buffers_computed = 0;

    for (cgltf_size mesh_idx = 0; mesh_idx < data->meshes_count; ++mesh_idx) {
      const cgltf_mesh* mesh = &data->meshes[mesh_idx];

      for (cgltf_size prim_idx = 0; prim_idx < mesh->primitives_count;
           ++prim_idx) {
        process_primitive(&options, data, &mesh->primitives[prim_idx],
                          *tangent_buffers, &tangent_buffers_computed);
      }
    }

    assert(tangent_buffers_computed == *num_tangent_buffers);
  }

  return cgltf_result_success;
}