//
// Created by martin on 25.03.22.
//

#ifndef ANIMTESTBED_ANIMGRAPHDATA_H
#define ANIMTESTBED_ANIMGRAPHDATA_H

#include <ozz/base/maths/soa_transform.h>

#include <cstring>
#include <iostream>
#include <list>
#include <map>
#include <string>
#include <vector>

#include "SyncTrack.h"
#include "ozz/animation/runtime/animation.h"
#include "ozz/animation/runtime/skeleton.h"
#include "ozz/base/containers/vector.h"

//
// Data types
//
struct AnimGraph;
struct AnimNode;

struct AnimData {
  ozz::vector<ozz::math::SoaTransform> m_local_matrices;
};

struct AnimDataRef {
  AnimData* ptr = nullptr;
};

struct AnimDataAllocator {
  struct AnimDataList {
    AnimData* m_anim_data = nullptr;
    AnimDataList* next = nullptr;
  };

  std::list<AnimData*> m_anim_data_list;
  size_t m_num_allocations = 0;

  ~AnimDataAllocator() {
    while (!m_anim_data_list.empty()) {
      AnimData* front = m_anim_data_list.front();
#ifdef ANIM_DATA_ALLOCATOR_DEBUG
      std::cout << "about to delete with size "
                << front->m_anim_data->m_local_matrices.size()
                << front->m_anim_data << std::endl;
#endif
      delete front;
      m_anim_data_list.pop_front();
    }
  }

  AnimData* allocate(ozz::animation::Skeleton* skeleton) {
    if (m_anim_data_list.empty()) {
      AnimData* result = new AnimData();
      result->m_local_matrices.resize(skeleton->num_soa_joints());
#ifdef ANIM_DATA_ALLOCATOR_DEBUG
      std::cout << "Allocated with size " << result->m_local_matrices.size()
                << " " << result << std::endl;
#endif
      m_num_allocations++;
      return result;
    }

    AnimData* result = m_anim_data_list.front();
    m_anim_data_list.pop_front();

#ifdef ANIM_DATA_ALLOCATOR_DEBUG
    std::cout << "Reusing buffer with size " << result->m_local_matrices.size()
              << " " << result << std::endl;
#endif

    return result;
  }

  void free(AnimData* anim_data) {
#ifdef ANIM_DATA_ALLOCATOR_DEBUG
    std::cout << "Storing buffer with size "
              << anim_data->m_local_matrices.size() << " " << anim_data
              << std::endl;
#endif

    m_anim_data_list.push_front(anim_data);
  }

  size_t size() { return m_anim_data_list.size(); }
};

struct AnimGraphContext {
  AnimGraph* m_graph = nullptr;
  ozz::animation::Skeleton* m_skeleton = nullptr;

  typedef std::map<std::string, ozz::animation::Animation*> AnimationFileMap;
  AnimationFileMap m_animation_map;

  void freeAnimations() {
    AnimationFileMap::iterator animation_map_iter = m_animation_map.begin();

    while (animation_map_iter != m_animation_map.end()) {
      delete animation_map_iter->second;
      animation_map_iter++;
    }
  }
};

union Vec3 {
  struct {
    float x;
    float y;
    float z;
  };

  float v[3] = {0};
};

union Quat {
  struct {
    float x;
    float y;
    float z;
    float w;
  };

  float v[4] = {0};
};

enum class SocketType {
  SocketTypeUndefined = 0,
  SocketTypeBool,
  SocketTypeAnimation,
  SocketTypeInt,
  SocketTypeFloat,
  SocketTypeVec3,
  SocketTypeQuat,
  SocketTypeString,
  SocketTypeLast
};

constexpr size_t cSocketStringValueMaxLength = 256;

static const char* SocketTypeNames[] =
    {"", "Bool", "Animation", "Int", "Float", "Vec3", "Quat", "String"};

enum SocketFlags { SocketFlagNone = 0, SocketFlagAffectsTime = 1 };

struct Socket {
  std::string m_name;
  SocketType m_type = SocketType::SocketTypeUndefined;
  union SocketValue {
    bool flag;
    int int_value;
    float float_value;
    Vec3 vec3;
    Quat quat;
  };
  SocketValue m_value = {0};
  std::string m_value_string;
  union SocketReference {
    void* ptr;
    void** ptr_ptr;
  };
  SocketReference m_reference = {0};
  SocketFlags m_flags = SocketFlagNone;
  size_t m_type_size = 0;

  template <typename T>
  void SetValue(const T value) {
    if constexpr (std::is_same<T, bool>::value) {
      m_value.flag = value;
    }

    if constexpr (std::is_same<T, int>::value) {
      m_value.int_value = value;
    }

    if constexpr (std::is_same<T, float>::value) {
      m_value.float_value = value;
    }

    if constexpr (std::is_same<T, Vec3>::value) {
      m_value.vec3 = value;
    }

    if constexpr (std::is_same<T, Quat>::value) {
      m_value.quat = value;
    }

    if constexpr (std::is_same<T, std::string>::value) {
      m_value_string = value;
    }
  }

  template <typename T>
  T GetValue() const {
    if constexpr (std::is_same<T, bool>::value) {
      return m_value.flag;
    }

    if constexpr (std::is_same<T, int>::value) {
      return m_value.int_value;
    }

    if constexpr (std::is_same<T, float>::value) {
      return m_value.float_value;
    }

    if constexpr (std::is_same<T, Vec3>::value) {
      return m_value.vec3;
    }

    if constexpr (std::is_same<T, Quat>::value) {
      return m_value.quat;
    }

    if constexpr (std::is_same<T, std::string>::value) {
      return m_value_string;
    }

    return T();
  }
};

template <typename T>
SocketType GetSocketType() {
  if constexpr (std::is_same<T, bool>::value) {
    return SocketType::SocketTypeBool;
  }

  if constexpr (std::is_same<T, AnimData>::value) {
    return SocketType::SocketTypeAnimation;
  }

  if constexpr (std::is_same<T, int>::value) {
    return SocketType::SocketTypeInt;
  }

  if constexpr (std::is_same<T, float>::value) {
    return SocketType::SocketTypeFloat;
  }

  if constexpr (std::is_same<T, Vec3>::value) {
    return SocketType::SocketTypeVec3;
  }

  if constexpr (std::is_same<T, Quat>::value) {
    return SocketType::SocketTypeQuat;
  }

  if constexpr (std::is_same<T, std::string>::value) {
    return SocketType::SocketTypeString;
  }

  assert(false && "This should not be reachable");
  abort();
  return SocketType::SocketTypeUndefined;
}

struct AnimGraphConnection {
  AnimNode* m_source_node = nullptr;
  Socket m_source_socket;
  AnimNode* m_target_node = nullptr;
  Socket m_target_socket;
  bool m_crosses_hierarchy = false;
};


struct NodeDescriptorBase {
  std::vector<Socket> m_inputs;
  std::vector<Socket> m_outputs;
  std::vector<Socket> m_properties;

  template <typename T>
  bool RegisterInput(
      const char* name,
      T** value_ptr_ptr,
      SocketFlags flags = SocketFlags::SocketFlagNone) {
    return RegisterSocket(name, value_ptr_ptr, m_inputs, flags);
  }

  template <typename T>
  bool RegisterOutput(
      const char* name,
      T** value_ptr_ptr,
      SocketFlags flags = SocketFlags::SocketFlagNone) {
    return RegisterSocket(name, value_ptr_ptr, m_outputs, flags);
  }

  template <typename T>
  bool RegisterProperty(
      const char* name,
      T* value_ptr,
      SocketFlags flags = SocketFlags::SocketFlagNone) {
    for (int i = 0; i < m_properties.size(); i++) {
      if (m_properties[i].m_name == name) {
        return false;
      }
    }

    Socket socket;
    socket.m_name = name;
    socket.m_type = GetSocketType<T>();
    socket.m_reference.ptr = static_cast<void*>(value_ptr);
    socket.m_flags = flags;
    socket.m_type_size = sizeof(T);

    m_properties.push_back(socket);

    return true;
  }

  template <typename T>
  T* GetInput(const char* name) {
    Socket* socket = FindSocket(name, m_inputs);
    assert(GetSocketType<T>() == socket->m_type);
    return *socket->m_reference.ptr_ptr;
  }

  template <typename T>
  T GetInputValue(const char* name) const {
    const Socket* socket = FindSocket(name, m_inputs);
    assert(GetSocketType<T>() == socket->m_type);
    return socket->GetValue<T>();
  }

  template <typename T>
  void SetInput(const char* name, T* value_ptr) {
    Socket* socket = FindSocket(name, m_inputs);
    assert(GetSocketType<T>() == socket->m_type);
    *socket->m_reference.ptr_ptr = value_ptr;
  }

  template <typename T>
  void SetInputValue(const char* name, T value) {
    Socket* socket = FindSocket(name, m_inputs);
    assert(GetSocketType<T>() == socket->m_type);
    socket->SetValue(value);
  }

  void SetInputUnchecked(const char* name, void* value_ptr) {
    Socket* socket = FindSocket(name, m_inputs);
    *socket->m_reference.ptr_ptr = value_ptr;
  }

  Socket* GetInputSocket(const char* name) {
    return FindSocket(name, m_inputs);
  }

  int GetInputIndex(const char* name) {
    return FindSocketIndex(name, m_inputs);
  }

  template <typename T>
  void SetOutput(const char* name, T* value_ptr) {
    Socket* socket = FindSocket(name, m_outputs);
    assert(GetSocketType<T>() == socket->m_type);
    *socket->m_reference.ptr_ptr = value_ptr;
  }

  void SetOutputUnchecked(const char* name, void* value_ptr) {
    Socket* socket = FindSocket(name, m_outputs);
    *socket->m_reference.ptr_ptr = value_ptr;
  }

  Socket* GetOutputSocket(const char* name) {
    return FindSocket(name, m_outputs);
  }

  int GetOutputIndex(const char* name) {
    return FindSocketIndex(name, m_outputs);
  }

  /** Sets value of an AnimNode Socket.
   *
   * @note Should only be used when the NodeDescriptor is associated with an AnimNode instance.
   *
   * @tparam T can be any AnimGraph data type.
   * @param Socket name
   * @param value
   */
  template <typename T>
  void SetProperty(const char* name, const T& value) {
    Socket* socket = FindSocket(name, m_properties);
    assert(GetSocketType<T>() == socket->m_type);
    *static_cast<T*>(socket->m_reference.ptr) = value;
  }

  /** Sets value of an AnimNodeResource Socket.
   *
   * @note Should only be used when the NodeDescriptor is associated with an AnimNodeResource instance. For AnimNode instances use Socket::SetProperty().
   *
   * @tparam T can be any AnimGraph data type.
   * @param Socket name
   * @param value
   */
  template <typename T>
  void SetPropertyValue(const char* name, const T& value) {
    Socket* socket = FindSocket(name, m_properties);
    assert(GetSocketType<T>() == socket->m_type);
    socket->SetValue(value);
  }

  template <typename T>
  const T& GetProperty(const char* name) {
    Socket* socket = FindSocket(name, m_properties);
    assert(GetSocketType<T>() == socket->m_type);
    return *static_cast<T*>(socket->m_reference.ptr);
  }

  template <typename T>
  T GetPropertyValue(const char* name) {
    Socket* socket = FindSocket(name, m_properties);
    assert(GetSocketType<T>() == socket->m_type);
    return socket->GetValue<T>();
  }

  virtual void UpdateFlags(){};

 protected:
  Socket* FindSocket(const char* name, std::vector<Socket>& sockets) {
    for (int i = 0, n = sockets.size(); i < n; i++) {
      if (sockets[i].m_name == name) {
        return &sockets[i];
      }
    }

    return nullptr;
  }

  const Socket* FindSocket(const char* name, const std::vector<Socket>& sockets) const {
    for (int i = 0, n = sockets.size(); i < n; i++) {
      if (sockets[i].m_name == name) {
        return &sockets[i];
      }
    }

    return nullptr;
  }

  int FindSocketIndex(const char* name, std::vector<Socket>& sockets) {
    for (int i = 0, n = sockets.size(); i < n; i++) {
      if (sockets[i].m_name == name) {
        return i;
      }
    }

    return -1;
  }

  template <typename T>
  bool RegisterSocket(
      const char* name,
      T** value_ptr_ptr,
      std::vector<Socket>& sockets,
      SocketFlags flags) {
    for (int i = 0; i < sockets.size(); i++) {
      if (sockets[i].m_name == name) {
        return false;
      }
    }

    Socket socket;
    socket.m_name = name;
    socket.m_type = GetSocketType<T>();
    socket.m_reference.ptr_ptr = (void**)(value_ptr_ptr);
    socket.m_type_size = sizeof(T);
    socket.m_flags = flags;

    sockets.push_back(socket);

    return true;
  }
};

template <typename T>
struct NodeDescriptor : public NodeDescriptorBase {
  virtual ~NodeDescriptor() {}
};

struct AnimNode;

template <typename T>
NodeDescriptorBase* CreateNodeDescriptor(AnimNode* node) {
  return new NodeDescriptor<T>(dynamic_cast<T*>(node));
}

#endif  //ANIMTESTBED_ANIMGRAPHDATA_H