//
// 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 NodeSocketAccessor<BlendTreeNode> : public NodeSocketAccessorBase {
  NodeSocketAccessor(AnimNode* 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 NodeSocketAccessor<Blend2Node> : public NodeSocketAccessorBase {
  NodeSocketAccessor(AnimNode* node_) {
    Blend2Node* node = dynamic_cast<Blend2Node*>(node_);
    RegisterInput("Input0", &node->i_input0);
    RegisterInput("Input1", &node->i_input1);
    RegisterInput(
        "Weight",
        &node->i_blend_weight,
        SocketFlags::SocketFlagAffectsTime);

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

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

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

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

//
// 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 = time_last;
    m_time_now = 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 NodeSocketAccessor<SpeedScaleNode> : public NodeSocketAccessorBase {
  NodeSocketAccessor(AnimNode* node_) {
    SpeedScaleNode* node = dynamic_cast<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::SamplingCache m_sampling_cache;
  ozz::animation::Animation* m_animation = nullptr;

  virtual ~AnimSamplerNode();
  virtual bool Init(AnimGraphContext& context) override;
  virtual void Evaluate(AnimGraphContext& context) override;
};

template <>
struct NodeSocketAccessor<AnimSamplerNode> : public NodeSocketAccessorBase {
  NodeSocketAccessor(AnimNode* node_) {
    AnimSamplerNode* node = dynamic_cast<AnimSamplerNode*>(node_);
    RegisterOutput("Output", &node->o_output);

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

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

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

    o_output = *i_input0 + *i_input1;
  }
};

template <>
struct NodeSocketAccessor<MathAddNode> : public NodeSocketAccessorBase {
  NodeSocketAccessor(AnimNode* node_) {
    MathAddNode* node = dynamic_cast<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 = {0.f, 0.f, 0.f};

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

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

template <>
struct NodeSocketAccessor<MathFloatToVec3Node> : public NodeSocketAccessorBase {
  NodeSocketAccessor(AnimNode* node_) {
    MathFloatToVec3Node* node = dynamic_cast<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 == "BlendTree") {
    result = new BlendTreeNode;
  } else if (name == "MathAddNode") {
    result = new MathAddNode;
  } else if (name == "MathFloatToVec3Node") {
    result = new MathFloatToVec3Node;
  }

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

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

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

#endif  //ANIMTESTBED_ANIMGRAPHNODES_H