#include "animation.h"
#include "lib/array.h"
#include "lib/assert.h"
#include "lib/math.h"
#include "lib/typedefs.h"
#include "renderer.h"
#include "vendor/cglm/types.h"

void apply_joint_transform_recursive(Joint_Array joints,
                                     mat4_Array joint_matrices, u32 joint_idx,
                                     mat4 parent_transform) {
  mat4 *joint_transform = arr_get_ptr(joint_matrices, joint_idx);
  mat4_mul(parent_transform, *joint_transform, *joint_transform);
  Joint *joint = arr_get_ptr(joints, joint_idx);
  arr_foreach(joint->children, u32, child_idx) {
    apply_joint_transform_recursive(joints, joint_matrices, child_idx,
                                    *joint_transform);
  }

  mat4_mul(*joint_transform, joint->inverse_bind_matrix, *joint_transform);
}

b32 find_start_end_keyframe(AnimationState *animation,
                            _out_ i32 *start_keyframe_idx,
                            _out_ i32 *end_keyframe_idx, _out_ f32 *t) {
  debug_assert(animation->animation->keyframes.len > 0);
  if (animation->animation->keyframes.len <= 0) {
    return false;
  }

  u32 len = animation->animation->keyframes.len;
  f32 time = animation->time;

  // handle wrap case
  if (animation->animation->keyframes.items[0].timestamp >= time) {
    *start_keyframe_idx = len - 1;
    *end_keyframe_idx = 0;
    *t = lerp_inverse(
        animation->animation->keyframes.items[*start_keyframe_idx].timestamp,
        animation->animation->keyframes.items[*end_keyframe_idx].timestamp,
        animation->time);
    return true;
  }

  // binary search for the keyframe
  u32 left = 0;
  u32 right = len - 1;

  while (left < right) {
    u32 mid = left + (right - left) / 2;
    if (animation->animation->keyframes.items[mid].timestamp <= time) {
      left = mid + 1;
    } else {
      right = mid;
    }
  }

  // left is now the first keyframe with timestamp > time
  *start_keyframe_idx = left - 1;
  *end_keyframe_idx = left;
  *t = lerp_inverse(
      animation->animation->keyframes.items[*start_keyframe_idx].timestamp,
      animation->animation->keyframes.items[*end_keyframe_idx].timestamp,
      animation->time);

  return true;
}

void animation_update(AnimationState *animation, f32 dt) {
  f32 length = animation->animation->length;

  animation->time += dt * animation->speed;

  if (length > 0.0f) {
    // Use fmod to handle multiple wraps and negative values
    animation->time = fmodf(animation->time, length);

    // Ensure positive time for negative wrap-around
    if (animation->time < 0.0f) {
      animation->time += length;
    }
  }
}

void animation_evaluate(AnimationState *animation,
                        _out_ mat4_Array joint_matrices) {
  assert(animation && animation->animation);
  if (!animation || !animation->animation) {
    return;
  }

  i32 start_keyframe_idx, end_keyframe_idx;
  f32 percent;
  if (!find_start_end_keyframe(animation, &start;_keyframe_idx,
                               &end;_keyframe_idx, &percent;)) {
    return;
  }

  Keyframe *start_keyframe =
      &animation;->animation->keyframes.items[start_keyframe_idx];
  Keyframe *end_keyframe =
      &animation;->animation->keyframes.items[end_keyframe_idx];

  // evaluate all joints for keyframe and store local transforms
  for (u32 i = 0; i < start_keyframe->joint_transforms.len; i++) {
    KeyframeJointTransform *start_joint_transform =
        &start;_keyframe->joint_transforms.items[i];
    KeyframeJointTransform *end_joint_transform =
        &end;_keyframe->joint_transforms.items[i];

    // again, every keyframe defines every joint in the same order
    assert(start_joint_transform->index == end_joint_transform->index);

    vec3 translation;
    quaternion rotation;
    vec3_lerp(start_joint_transform->translation,
              end_joint_transform->translation, percent, translation);
    quat_slerp(start_joint_transform->rotation, end_joint_transform->rotation,
               percent, rotation);

    mat4 *joint_mat = arr_get_ptr(joint_matrices, start_joint_transform->index);
    debug_assert(joint_mat);
    if (!joint_mat) {
      continue;
    }

    mat_tr(translation, rotation, *joint_mat);
  }

  // apply hierarchy
  Model3DData *model = animation->animation->model;
  Joint_Array model_joints =
      arr_from_c_array(Joint, model->joints, model->len_joints);
  apply_joint_transform_recursive(model_joints, joint_matrices, 0,
                                  MAT4_IDENTITY);
}