#include "animation_impl.h"

#include "node_impl.h"
#include "scene_memory.h"

#include <string.h>

// #include <log/log.h> // Debugging.

static const R PLAYBACK_UNINITIALIZED = -1;

static joint_idx get_anima_root_joint_index(Anima* anima) {
  assert(anima);
  assert(anima->num_joints > 0);
  assert(anima->num_joints < GFX_MAX_NUM_JOINTS);
  return anima->num_joints - 1;
}

static Joint* get_anima_root_joint(Anima* anima) {
  assert(anima);
  return &anima->joints[get_anima_root_joint_index(anima)];
}

static const Joint* get_anima_joint(const Anima* anima, joint_idx index) {
  assert(anima);
  assert(index < GFX_MAX_NUM_JOINTS);
  assert(index != INDEX_NONE);
  assert(index < anima->num_joints);
  return &anima->joints[index];
}

static Joint* get_anima_joint_mut(Anima* anima, joint_idx index) {
  return (Joint*)get_anima_joint(anima, index);
}

static const Joint* get_skeleton_joint(
    const Anima* anima, const Skeleton* skeleton, joint_idx index) {
  assert(anima);
  assert(skeleton);
  return get_anima_joint(anima, skeleton->joints[index]);
}

static void set_joint_parent(
    Anima* anima, joint_idx joint_index, joint_idx parent_index) {
  assert(anima);
  assert(joint_index != INDEX_NONE);
  assert(joint_index != get_anima_root_joint_index(anima));
  assert(parent_index != INDEX_NONE);

  Joint* parent = get_anima_joint_mut(anima, parent_index);

  if (parent->child == INDEX_NONE) {
    parent->child = joint_index;
  } else {
    // Find the last child in the chain of children.
    Joint* child = get_anima_joint_mut(anima, parent->child);
    while (child->next != INDEX_NONE) {
      child = get_anima_joint_mut(anima, child->next);
    }
    // Wire up this joint as the last child's sibling.
    child->next = joint_index;
  }
}

static void make_joint(Anima* anima, const JointDesc* desc, Joint* joint) {
  assert(anima);
  assert(desc);
  assert(joint);

  // The joint matrix needs to be initialized so that meshes look right even if
  // no animation is played. Initializing joint matrices to the identity makes
  // meshes appear in their bind pose.
  joint->child           = INDEX_NONE;
  joint->next            = INDEX_NONE;
  joint->transform       = mat4_id();
  joint->inv_bind_matrix = desc->inv_bind_matrix;
  joint->joint_matrix    = mat4_id();
  joint->box             = desc->box;
}

static Skeleton* make_skeleton(const SkeletonDesc* desc) {
  assert(desc);
  assert(desc->num_joints <= GFX_MAX_NUM_JOINTS);

  Skeleton* skeleton   = mem_alloc_skeleton();
  skeleton->num_joints = desc->num_joints;
  memcpy(
      skeleton->joints, desc->joints,
      desc->num_joints * sizeof(skeleton->joints[0]));
  return skeleton;
}

static Animation* make_animation(const AnimationDesc* desc) {
  assert(desc);
  assert(desc->num_channels < GFX_MAX_NUM_CHANNELS);

  Animation* animation    = mem_alloc_animation();
  animation->name         = desc->name;
  animation->duration     = 0;
  animation->num_channels = desc->num_channels;
  R start_time            = 0;
  R end_time              = 0;

  for (size_t c = 0; c < desc->num_channels; ++c) {
    const ChannelDesc* channel_desc = &desc->channels[c];
    Channel*           channel      = &animation->channels[c];

    channel->target        = channel_desc->target;
    channel->type          = channel_desc->type;
    channel->interpolation = channel_desc->interpolation;
    channel->num_keyframes = channel_desc->num_keyframes;
    assert(channel_desc->num_keyframes < GFX_MAX_NUM_KEYFRAMES);

    for (size_t k = 0; k < channel_desc->num_keyframes; ++k) {
      const KeyframeDesc* keyframe_desc = &channel_desc->keyframes[k];
      Keyframe*           keyframe      = &channel->keyframes[k];

      keyframe->time        = keyframe_desc->time;
      keyframe->translation = keyframe_desc->translation;
      keyframe->rotation    = keyframe_desc->rotation;

      start_time = keyframe->time < start_time ? keyframe->time : start_time;
      end_time   = keyframe->time > end_time ? keyframe->time : end_time;
    }
  }

  // LOGD("Animation start/end: %f / %f", start_time, end_time);
  animation->duration = end_time - start_time;
  assert(animation->duration >= 0);
  return animation;
}

Anima* gfx_make_anima(const AnimaDesc* desc) {
  assert(desc);
  assert(desc->num_joints > 0);
  assert(desc->num_joints <= GFX_MAX_NUM_JOINTS);
  // All joints should have a parent except for the root.
  for (size_t i = 0; i < desc->num_joints - 1; ++i) {
    const joint_idx parent = desc->joints[i].parent;
    assert(parent != INDEX_NONE);
    assert(parent < desc->num_joints);
  }
  // The root should have no parent.
  assert(desc->joints[desc->num_joints - 1].parent == INDEX_NONE);

  Anima* anima = mem_alloc_anima();

  // Wire the skeletons in the same order they are given in the descriptor.
  Skeleton* last_skeleton = 0;
  for (size_t i = 0; i < desc->num_skeletons; ++i) {
    Skeleton*          skeleton       = make_skeleton(&desc->skeletons[i]);
    const skeleton_idx skeleton_index = mem_get_skeleton_index(skeleton);
    if (last_skeleton == 0) {
      anima->skeleton = skeleton_index;
    } else {
      last_skeleton->next = skeleton_index;
    }
    last_skeleton = skeleton;
  }

  // Wire the animations in the same order they are given in the descriptor.
  Animation* last_animation = 0;
  for (size_t i = 0; i < desc->num_animations; ++i) {
    Animation*          animation       = make_animation(&desc->animations[i]);
    const animation_idx animation_index = mem_get_animation_index(animation);
    if (last_animation == 0) {
      anima->animation = animation_index;
    } else {
      last_animation->next = animation_index;
    }
    last_animation = animation;
  }

  // Create joints.
  anima->num_joints = desc->num_joints;
  // Initialize all joints.
  // Child and sibling pointers must be initialized before wiring up the
  // hierarchy.
  for (size_t i = 0; i < desc->num_joints; ++i) {
    Joint* joint = get_anima_joint_mut(anima, i);
    make_joint(anima, &desc->joints[i], joint);
  }
  // Wire up joints to their parents. -1 to skip the root.
  for (size_t i = 0; i < desc->num_joints - 1; ++i) {
    set_joint_parent(anima, i, desc->joints[i].parent);
  }

  return anima;
}

void gfx_destroy_anima(Anima** anima) {
  assert(anima);

  if (*anima) {
    for (skeleton_idx i = (*anima)->skeleton; i.val != 0;) {
      Skeleton* skeleton = mem_get_skeleton(i);
      i                  = skeleton->next;
      mem_free_skeleton(&skeleton);
    }

    for (animation_idx i = (*anima)->animation; i.val != 0;) {
      Animation* animation = mem_get_animation(i);
      i                    = animation->next;
      mem_free_animation(&animation);
    }

    if ((*anima)->parent.val) {
      gfx_del_node((*anima)->parent);
    }

    mem_free_anima(anima);
  }
}

static Animation* find_animation(animation_idx index, const char* name) {
  assert(name);

  while (index.val != 0) {
    Animation* animation = mem_get_animation(index);
    if (sstring_eq_cstr(animation->name, name)) {
      // LOGD(
      //     "Found animation at index %u, %s - %s", index.val,
      //     sstring_cstr(&animation->name), name);
      // LOGD("Animation has duration %f", animation->duration);
      return animation;
    }
    index = animation->next;
  }

  return 0;
}

bool gfx_play_animation(Anima* anima, const AnimationPlaySettings* settings) {
  assert(anima);
  assert(settings);

  // TODO: Should we animate at t=0 here to kickstart the animation? Otherwise
  //  the client is forced to call gfx_update_animation() to do this.
  Animation* animation = find_animation(anima->animation, settings->name);
  if (!animation) {
    return false;
  }
  // Playback initialized on first call to update().
  AnimationState* state = &anima->state;
  state->start_time     = PLAYBACK_UNINITIALIZED;
  state->animation      = mem_get_animation_index(animation);
  state->loop           = settings->loop;
  return true;
}

static void gfx_set_joint_position(Joint* joint, vec3 position) {
  assert(joint);
  mat4_set_v3(&joint->transform, position);
}

static void gfx_set_joint_rotation(Joint* joint, quat rotation) {
  assert(joint);
  mat4_set_3x3(&joint->transform, mat4_from_quat(rotation));
}

static void find_keyframes(const Channel* channel, R t, int* prev, int* next) {
  assert(channel);
  assert(prev);
  assert(next);

  *prev = -1;
  *next = 0;
  while (((*next + 1) < (int)channel->num_keyframes) &&
         (t >= channel->keyframes[*next + 1].time)) {
    (*prev)++;
    (*next)++;
  }
}

static R normalize_time(R a, R b, R t) {
  assert(a <= t);
  assert(t <= b);
  return (t - a) / (b - a);
}

static quat interpolate_rotation(
    const Channel* channel, int prev, int next, R t) {
  assert(channel);

  if (next == 0) {
    // Animation has not started at this point in time yet.
    return channel->keyframes[next].rotation;
  } else {
    switch (channel->interpolation) {
    case StepInterpolation:
      return channel->keyframes[prev].rotation;
    case LinearInterpolation: {
      const R normalized_t = normalize_time(
          channel->keyframes[prev].time, channel->keyframes[next].time, t);
      return qnormalize(qslerp(
          channel->keyframes[prev].rotation, channel->keyframes[next].rotation,
          normalized_t));
      break;
    }
    case CubicSplineInterpolation:
      assert(false); // TODO
      return qmake(0, 0, 0, 0);
    default:
      assert(false);
      return qmake(0, 0, 0, 0);
    }
  }
}

static vec3 interpolate_translation(
    const Channel* channel, int prev, int next, R t) {
  assert(channel);

  if (next == 0) {
    // Animation has not started at this point in time yet.
    return channel->keyframes[next].translation;
  } else {
    switch (channel->interpolation) {
    case StepInterpolation:
      return channel->keyframes[prev].translation;
    case LinearInterpolation: {
      const R normalized_t = normalize_time(
          channel->keyframes[prev].time, channel->keyframes[next].time, t);
      return vec3_lerp(
          channel->keyframes[prev].translation,
          channel->keyframes[next].translation, normalized_t);
      break;
    }
    case CubicSplineInterpolation:
      assert(false); // TODO
      return vec3_make(0, 0, 0);
    default:
      assert(false);
      return vec3_make(0, 0, 0);
    }
  }
}

static void animate_channel(Anima* anima, const Channel* channel, R t) {
  assert(anima);
  assert(channel);
  assert(channel->target < anima->num_joints);

  int prev, next;
  find_keyframes(channel, t, &prev, &next);

  // Note that not all channels extend to the duration of an animation; some
  // channels may stop animating their targets earlier. Clamp the animation time
  // to the channel's end keyframe to make the rest of the math (normalize_time)
  // work.
  t = t > channel->keyframes[next].time ? channel->keyframes[next].time : t;

  Joint* target = get_anima_joint_mut(anima, channel->target);

  switch (channel->type) {
  case RotationChannel: {
    const quat rotation = interpolate_rotation(channel, prev, next, t);
    gfx_set_joint_rotation(target, rotation);
    break;
  }
  case TranslationChannel: {
    const vec3 translation = interpolate_translation(channel, prev, next, t);
    gfx_set_joint_position(target, translation);
    break;
  }
  // Not yet supported.
  case ScaleChannel:
  case WeightsChannel:
  default:
    // TODO: Add back the assertion or add support for scaling.
    // assert(false);
    break;
  }
}

static void compute_joint_matrices_rec(
    Anima* anima, Joint* joint, const mat4* parent_global_joint_transform,
    const mat4* root_inv_global_transform) {
  assert(anima);
  assert(joint);
  assert(parent_global_joint_transform);
  assert(root_inv_global_transform);

  const mat4 global_joint_transform =
      mat4_mul(*parent_global_joint_transform, joint->transform);

  // Compute this joint's matrix.
  joint->joint_matrix = mat4_mul(
      *root_inv_global_transform,
      mat4_mul(global_joint_transform, joint->inv_bind_matrix));

  // Recursively compute the joint matrices for this joint's siblings.
  if (joint->next != INDEX_NONE) {
    Joint* sibling = get_anima_joint_mut(anima, joint->next);

    compute_joint_matrices_rec(
        anima, sibling, parent_global_joint_transform,
        root_inv_global_transform);
  }

  // Recursively compute the joint matrices for this joint's children.
  if (joint->child != INDEX_NONE) {
    Joint* child = get_anima_joint_mut(anima, joint->child);

    compute_joint_matrices_rec(
        anima, child, &global_joint_transform, root_inv_global_transform);
  }
}

void gfx_update_animation(Anima* anima, R t) {
  assert(anima);

  AnimationState* state = &anima->state;
  if (state->animation.val == 0) {
    return; // No active animation.
  }
  const Animation* animation = mem_get_animation(state->animation);
  assert(animation);

  // On a call to play(), the start time is set to -1 to signal that the
  // animation playback has not yet been initialized.
  if (state->start_time == PLAYBACK_UNINITIALIZED) {
    state->start_time = t;
  }
  // Locate the current time point inside the animation's timeline.
  assert(t >= state->start_time);
  assert(animation->duration >= 0.0);
  const R local_time     = t - state->start_time;
  const R animation_time = state->loop
                               ? rmod(local_time, animation->duration)
                               : clamp(local_time, 0.0, animation->duration);

  // LOGD(
  //     "animation_time = %f, animation duration: %f", animation_time,
  //     animation->duration);

  // Play through the animation to transform skeleton nodes.
  for (size_t i = 0; i < animation->num_channels; ++i) {
    const Channel* channel = &animation->channels[i];
    animate_channel(anima, channel, animation_time);
  }

  // Compute joint matrices after having transformed the skeletons.
  //
  // The anima's parent node is the common ancestor of all skeletons, and its
  // transform maps the skeletons from object space to world space. This is the
  // transform used as the "global transform" in the joint matrix equations.
  //
  // Joint matrix calculation begins by descending from the anima's root joint,
  // which we have constructed to be the common root of all skeletons.
  //
  // This procedure touches every joint exactly once.
  SceneNode* root_node = mem_get_node(anima->parent);
  // LOGD("Root: %u, child: %u", anima->parent.val, root->child.val);
  const mat4 root_global_transform = gfx_get_node_global_transform(root_node);
  const mat4 root_inv_global_transform = mat4_inverse(root_global_transform);

  Joint* root_joint = get_anima_root_joint(anima);
  compute_joint_matrices_rec(
      anima, root_joint, &root_global_transform, &root_inv_global_transform);
}

const Skeleton* gfx_get_anima_skeleton(const Anima* anima, size_t i) {
  assert(anima);

  skeleton_idx    skeleton_index = anima->skeleton;
  const Skeleton* skeleton       = mem_get_skeleton(skeleton_index);

  for (size_t j = 1; j < i; ++j) {
    skeleton_index = skeleton->next;
    mem_get_skeleton(skeleton_index);
  }

  return skeleton;
}

size_t gfx_get_skeleton_num_joints(const Skeleton* skeleton) {
  assert(skeleton);
  return skeleton->num_joints;
}

bool gfx_joint_has_box(
    const Anima* anima, const Skeleton* skeleton, size_t joint_index) {
  assert(anima);
  assert(skeleton);
  assert(joint_index < skeleton->num_joints);

  const Joint* joint = get_skeleton_joint(anima, skeleton, joint_index);
  return !aabb3_is_empty(joint->box);
}

Box gfx_get_joint_box(
    const Anima* anima, const Skeleton* skeleton, size_t joint_index) {
  assert(anima);
  assert(skeleton);

  const Joint* joint = get_skeleton_joint(anima, skeleton, joint_index);

  // Transform the box to anima space.
  // Note that joint matrices do not usually have a translation since joints
  // mostly just rotate with respect to their parent.
  const vec3 pmin = joint->box.min;
  const vec3 pmax = joint->box.max;
  return (Box){
      .vertices = {
                   mat4_mul_vec3(
                   joint->joint_matrix, vec3_make(pmin.x, pmin.y, pmax.z), 1),
                   mat4_mul_vec3(
                   joint->joint_matrix, vec3_make(pmax.x, pmin.y, pmax.z), 1),
                   mat4_mul_vec3(
                   joint->joint_matrix, vec3_make(pmax.x, pmax.y, pmax.z), 1),
                   mat4_mul_vec3(
                   joint->joint_matrix, vec3_make(pmin.x, pmax.y, pmax.z), 1),
                   mat4_mul_vec3(
                   joint->joint_matrix, vec3_make(pmin.x, pmin.y, pmin.z), 1),
                   mat4_mul_vec3(
                   joint->joint_matrix, vec3_make(pmax.x, pmin.y, pmin.z), 1),
                   mat4_mul_vec3(
                   joint->joint_matrix, vec3_make(pmax.x, pmax.y, pmin.z), 1),
                   mat4_mul_vec3(
                   joint->joint_matrix, vec3_make(pmin.x, pmax.y, pmin.z), 1),
                   }
  };
}