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

#ifndef ANIMTESTBED_ANIMGRAPHNODES_H
#define ANIMTESTBED_ANIMGRAPHNODES_H

#include <vector>

#include "AnimGraphData.h"
#include "AnimNode.h"
#include "SyncTrack.h"
#include "ozz/animation/runtime/sampling_job.h"

struct AnimNode;

//
// Blend2Node
//
struct Blend2Node : public AnimNode {
  AnimData* i_input0 = nullptr;
  AnimData* i_input1 = nullptr;
  AnimData* o_output = nullptr;
  float* i_blend_weight = nullptr;
  bool m_sync_blend = false;

  void MarkActiveInputs(
      const std::vector<AnimGraphConnection>& input_connections) override {
    for (const auto& input : input_connections) {
      AnimNode* input_node = input.m_source_node;
      if (input_node == nullptr) {
        continue;
      }

      if (input.m_target_socket_name == "Input0" && *i_blend_weight < 0.999) {
        input_node->Activate(m_tick_number);
        continue;
      }

      if (input.m_target_socket_name == "Input1" && *i_blend_weight > 0.001) {
        input_node->Activate(m_tick_number);
        continue;
      }
    }
  }

  void Evaluate(AnimGraphContext& context) override;
};

template <>
struct NodeDescriptor<Blend2Node> : public NodeDescriptorBase {
  explicit NodeDescriptor(Blend2Node* node) {
    RegisterInput("Input0", &node->i_input0);
    RegisterInput("Input1", &node->i_input1);
    RegisterInput("Weight", &node->i_blend_weight);

    RegisterOutput("Output", &node->o_output);

    RegisterProperty("Sync", &node->m_sync_blend);
  }

  virtual ~NodeDescriptor() = default;

  void UpdateFlags() override {
    Socket* weight_input_socket = FindSocket("Weight", m_inputs);
    assert(weight_input_socket != nullptr);

    if (GetProperty<bool>("Sync")) {
      weight_input_socket->m_flags = SocketFlags::SocketFlagAffectsTime;
    } else {
      weight_input_socket->m_flags = SocketFlags::SocketFlagNone;
    }
  }
};

//
// SpeedScaleNode
//
struct SpeedScaleNode : public AnimNode {
  AnimData* i_input = nullptr;
  AnimData* o_output = nullptr;
  float* i_speed_scale = nullptr;

  void UpdateTime(float time_last, float time_now) override {
    m_time_last = m_time_now;
    m_time_now = m_time_last + (time_now - time_last) * (*i_speed_scale);
    m_state = AnimNodeEvalState::TimeUpdated;
  }

  void Evaluate(AnimGraphContext& context) override {
    assert(i_input != nullptr);
    assert(o_output != nullptr);

    o_output->m_local_matrices = i_input->m_local_matrices;
  };
};

template <>
struct NodeDescriptor<SpeedScaleNode> : public NodeDescriptorBase {
  explicit NodeDescriptor(SpeedScaleNode* node) {
    RegisterInput(
        "SpeedScale",
        &node->i_speed_scale,
        SocketFlags::SocketFlagAffectsTime);
    RegisterInput("Input", &node->i_input);

    RegisterOutput("Output", &node->o_output);
  }

  virtual ~NodeDescriptor() = default;
};

//
// AnimSamplerNode
//
struct AnimSamplerNode : public AnimNode {
  AnimData* o_output = nullptr;
  std::string m_filename;
  ozz::animation::SamplingJob::Context m_sampling_context;
  ozz::animation::Animation* m_animation = nullptr;

  ~AnimSamplerNode() noexcept override;
  bool Init(AnimGraphContext& context) override;
  void UpdateTime(float time_last, float time_now) override {
    m_time_last = time_last;
    m_time_now = time_now;
    m_state = AnimNodeEvalState::TimeUpdated;
  }
  virtual void Evaluate(AnimGraphContext& context) override;
};

template <>
struct NodeDescriptor<AnimSamplerNode> : public NodeDescriptorBase {
  explicit NodeDescriptor(AnimSamplerNode* node) {
    RegisterOutput("Output", &node->o_output);

    RegisterProperty("Filename", &node->m_filename);
  }

  virtual ~NodeDescriptor() = default;
};

//
// LockTranslationNode
//
struct LockTranslationNode : public AnimNode {
  AnimData* i_input = nullptr;
  AnimData* o_output = nullptr;
  int m_locked_bone_index = 0;
  bool m_lock_x = false;
  bool m_lock_y = false;
  bool m_lock_z = false;

  virtual void Evaluate(AnimGraphContext& context) override;
};

template <>
struct NodeDescriptor<LockTranslationNode> : public NodeDescriptorBase {
  explicit NodeDescriptor(LockTranslationNode* node) {
    RegisterInput("Input", &node->i_input);
    RegisterOutput("Output", &node->o_output);

    RegisterProperty("BoneIndex", &node->m_locked_bone_index);
    RegisterProperty("LockAxisX", &node->m_lock_x);
    RegisterProperty("LockAxisY", &node->m_lock_y);
    RegisterProperty("LockAxisZ", &node->m_lock_z);
  }

  virtual ~NodeDescriptor() = default;
};

//
// ConstScalarNode
//
struct ConstScalarNode : public AnimNode {
  float* o_value = nullptr;
  float value = 0.f;

  virtual void Evaluate(AnimGraphContext& context) { *o_value = value; };
};

template <>
struct NodeDescriptor<ConstScalarNode> : public NodeDescriptorBase {
  explicit NodeDescriptor(ConstScalarNode* node) {
    RegisterOutput("ScalarOutput", &node->o_value);
    RegisterProperty("ScalarValue", &node->value);
  }

  virtual ~NodeDescriptor() = default;
};

//
// MathAddNode
//
struct MathAddNode : public AnimNode {
  float* i_input0 = nullptr;
  float* i_input1 = nullptr;
  float* o_output = nullptr;

  void Evaluate(AnimGraphContext& context) override {
    assert(i_input0 != nullptr);
    assert(i_input1 != nullptr);

    *o_output = *i_input0 + *i_input1;
  }
};

template <>
struct NodeDescriptor<MathAddNode> : public NodeDescriptorBase {
  explicit NodeDescriptor(MathAddNode* node) {
    RegisterInput("Input0", &node->i_input0);
    RegisterInput("Input1", &node->i_input1);
    RegisterOutput("Output", &node->o_output);
  }

  virtual ~NodeDescriptor() = default;
};

//
// MathFloatToVec3Node
//
struct MathFloatToVec3Node : public AnimNode {
  float* i_input0 = nullptr;
  float* i_input1 = nullptr;
  float* i_input2 = nullptr;
  Vec3* o_output = nullptr;

  void Evaluate(AnimGraphContext& context) override {
    assert(i_input0 != nullptr);
    assert(i_input1 != nullptr);
    assert(i_input2 != nullptr);

    o_output->v[0] = *i_input0;
    o_output->v[1] = *i_input1;
    o_output->v[2] = *i_input2;
  }
};

template <>
struct NodeDescriptor<MathFloatToVec3Node> : public NodeDescriptorBase {
  explicit NodeDescriptor(MathFloatToVec3Node* node) {
    RegisterInput("Input0", &node->i_input0);
    RegisterInput("Input1", &node->i_input1);
    RegisterInput("Input2", &node->i_input2);
    RegisterOutput("Output", &node->o_output);
  }

  virtual ~NodeDescriptor() = default;
};

AnimNode* AnimNodeFactory(const std::string& name);

NodeDescriptorBase* AnimNodeDescriptorFactory(
    const std::string& node_type_name,
    AnimNode* node);

NodeDescriptorBase* VirtualAnimNodeDescriptorFactory(
    const std::string& node_type_name);

#endif  //ANIMTESTBED_ANIMGRAPHNODES_H