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

#ifndef ANIMTESTBED_ANIMGRAPHNODES_H
#define ANIMTESTBED_ANIMGRAPHNODES_H

#include <vector>

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

struct AnimNode;

struct NodeInput {
  AnimNode* m_node;
  SocketType m_type = SocketType::SocketTypeUndefined;
  Socket* m_node_output_socket;
  std::string m_input_name;
};

enum class AnimNodeEvalState {
  Undefined,
  Deactivated,
  Activated,
  SyncTrackUpdated,
  TimeUpdated,
  Evaluated
};

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

struct AnimNode {
  std::string m_name;
  std::string m_node_type_name;
  float m_time_now = 0.f;
  float m_time_last = 0.f;
  size_t m_index = -1;
  AnimNodeEvalState m_state = AnimNodeEvalState::Undefined;
  SyncTrack m_sync_track;

  virtual ~AnimNode(){};

  virtual bool Init(AnimGraphContext& context) { return true; };

  virtual void MarkActiveInputs(const std::vector<AnimGraphConnection>& inputs) {
    for (size_t i = 0, n = inputs.size(); i < n; i++) {
      AnimNode* input_node = inputs[i].m_source_node;
      if (input_node != nullptr) {
        input_node->m_state = AnimNodeEvalState::Activated;
      }
    }
  }

  virtual void CalcSyncTrack(const std::vector<AnimGraphConnection>& inputs) {
    for (size_t i = 0, n = inputs.size(); i < n; i++) {
      AnimNode* input_node = inputs[i].m_source_node;
      if (input_node != nullptr
          && inputs[i].m_source_socket.m_type == SocketType::SocketTypeAnimation
          && input_node->m_state != AnimNodeEvalState::Deactivated) {
        m_sync_track = input_node->m_sync_track;
        return;
      }
    }
  }

  virtual void UpdateTime(float time_last, float time_now) {
    m_time_last = time_last;
    m_time_now = time_now;
    m_state = AnimNodeEvalState::TimeUpdated;
  }

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


//
// BlendTreeNode
//
struct BlendTreeNode : public AnimNode {};

template <>
struct NodeDescriptor<BlendTreeNode> : public NodeDescriptorBase {
  NodeDescriptor(BlendTreeNode* node_) {}
};

//
// 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;

  virtual void MarkActiveInputs(const std::vector<AnimGraphConnection>& inputs) override {
    for (size_t i = 0, n = inputs.size(); i < n; i++) {
      AnimNode* input_node = inputs[i].m_source_node;
      if (input_node == nullptr) {
        continue;
      }

      if (inputs[i].m_target_socket.m_name == "Input0" && *i_blend_weight < 0.999) {
        input_node->m_state = AnimNodeEvalState::Activated;
        continue;
      }

      if (inputs[i].m_target_socket.m_name == "Input1" && *i_blend_weight > 0.001) {
        input_node->m_state = AnimNodeEvalState::Activated;
        continue;
      }
    }
  }

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

template <>
struct NodeDescriptor<Blend2Node> : public NodeDescriptorBase {
  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);
  }

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

    if (GetProperty<bool>("Sync") == true) {
      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 {
  NodeDescriptor(SpeedScaleNode* node) {
    RegisterInput(
        "SpeedScale",
        &node->i_speed_scale,
        SocketFlags::SocketFlagAffectsTime);
    RegisterInput("Input", &node->i_input);

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


//
// 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;

  virtual ~AnimSamplerNode();
  virtual 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 {
  NodeDescriptor(AnimSamplerNode* node) {
    RegisterOutput("Output", &node->o_output);

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

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

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

template <>
struct NodeDescriptor<LockTranslationNode> : public NodeDescriptorBase {
  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);
  }
};


//
// 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 {
  NodeDescriptor(ConstScalarNode* node) {
    RegisterOutput("ScalarOutput", &node->o_value);
    RegisterProperty("ScalarValue", &node->value);
  }
};


//
// 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 {
  NodeDescriptor(MathAddNode* node) {
    RegisterInput("Input0", &node->i_input0);
    RegisterInput("Input1", &node->i_input1);
    RegisterOutput("Output", &node->o_output);
  }
};

//
// 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 {
  NodeDescriptor(MathFloatToVec3Node* node) {
    RegisterInput("Input0", &node->i_input0);
    RegisterInput("Input1", &node->i_input1);
    RegisterInput("Input2", &node->i_input2);
    RegisterOutput("Output", &node->o_output);
  }
};


static inline AnimNode* AnimNodeFactory(const std::string& name) {
  AnimNode* result;
  if (name == "Blend2") {
    result = new Blend2Node;
  } else if (name == "SpeedScale") {
    result = new SpeedScaleNode;
  } else if (name == "AnimSampler") {
    result = new AnimSamplerNode;
  } else if (name == "LockTranslationNode") {
    result = new LockTranslationNode;
  } else if (name == "BlendTree") {
    result = new BlendTreeNode;
  } else if (name == "MathAddNode") {
    result = new MathAddNode;
  } else if (name == "MathFloatToVec3Node") {
    result = new MathFloatToVec3Node;
  } else if (name == "ConstScalarNode") {
    result = new ConstScalarNode;
  }

  if (result != nullptr) {
    result->m_node_type_name = name;
    return result;
  }

  std::cerr << "Invalid node type: " << name << std::endl;
  return nullptr;
}

static inline NodeDescriptorBase* AnimNodeDescriptorFactory(
    const std::string& node_type_name,
    AnimNode* node) {
  if (node_type_name == "Blend2") {
    return CreateNodeDescriptor<Blend2Node>(node);
  } else if (node_type_name == "SpeedScale") {
    return CreateNodeDescriptor<SpeedScaleNode>(node);
  } else if (node_type_name == "AnimSampler") {
    return CreateNodeDescriptor<AnimSamplerNode>(node);
  } else if (node_type_name == "LockTranslationNode") {
    return CreateNodeDescriptor<LockTranslationNode>(node);
  } else if (node_type_name == "BlendTree") {
    return CreateNodeDescriptor<BlendTreeNode>(node);
  } else if (node_type_name == "MathAddNode") {
    return CreateNodeDescriptor<MathAddNode>(node);
  } else if (node_type_name == "MathFloatToVec3Node") {
    return CreateNodeDescriptor<MathFloatToVec3Node>(node);
  } else if (node_type_name == "ConstScalarNode") {
    return CreateNodeDescriptor<ConstScalarNode>(node);
  } else {
    std::cerr << "Invalid node type name " << node_type_name << "."
              << std::endl;
  }
  return nullptr;
}

#endif  //ANIMTESTBED_ANIMGRAPHNODES_H