AnimTestbed/src/AnimGraph/AnimGraphResource.h

//
// Created by martin on 17.03.24.
//

#ifndef ANIMTESTBED_ANIMGRAPHRESOURCE_H
#define ANIMTESTBED_ANIMGRAPHRESOURCE_H

#include "3rdparty/json/json.hpp"
#include "AnimGraphNodes.h"

struct AnimGraphBlendTree;
struct AnimGraphStateMachine;

struct AnimNodeResource {
  virtual ~AnimNodeResource() { delete m_virtual_socket_accessor; };

  std::string m_name;
  std::string m_node_type_name;
  NodeDescriptorBase* m_virtual_socket_accessor = nullptr;
  float m_position[2] = {0.f, 0.f};
};

static inline AnimNodeResource* AnimNodeResourceFactory(
    const std::string& node_type_name);

struct BlendTreeConnectionResource {
  int source_node_index = -1;
  std::string source_socket_name;
  int target_node_index = -1;
  std::string target_socket_name;

  bool operator==(const BlendTreeConnectionResource& other) const {
    return (
        source_node_index == other.source_node_index
        && target_node_index == other.target_node_index
        && source_socket_name == other.source_socket_name

        && target_socket_name == other.target_socket_name);
  }
};

struct BlendTreeResource {
  std::vector<std::vector<size_t> > m_node_input_connection_indices;
  std::vector<std::vector<size_t> > m_node_inputs_subtree;

  ~BlendTreeResource() { CleanupNodes(); }

  void Reset() {
    CleanupNodes();

    m_connections.clear();

    m_node_input_connection_indices.clear();
    m_node_inputs_subtree.clear();
  }

  void CleanupNodes() {
    for (AnimNodeResource* node_resource : m_nodes) {
      delete node_resource;
    }

    m_nodes.clear();
  }

  void InitGraphConnectors() {
    AddNode(AnimNodeResourceFactory("BlendTreeSockets"));
    AnimNodeResource* output_node = GetGraphOutputNode();
    output_node->m_name = "Outputs";
    output_node->m_position[0] = 200;

    AddNode(AnimNodeResourceFactory("BlendTreeSockets"));
    AnimNodeResource* input_node = GetGraphInputNode();
    input_node->m_name = "Inputs";
    input_node->m_position[0] = -200;
  }

  [[nodiscard]] AnimNodeResource* GetGraphOutputNode() const {
    return m_nodes[0];
  }
  [[nodiscard]] AnimNodeResource* GetGraphInputNode() const {
    return m_nodes[1];
  }
  Socket* GetGraphOutputSocket(const char* socket_name) const {
    return GetGraphOutputNode()->m_virtual_socket_accessor->GetInputSocket(
        socket_name);
  }
  Socket* GetGraphInputSocket(const char* socket_name) const {
    return GetGraphInputNode()->m_virtual_socket_accessor->GetOutputSocket(
        socket_name);
  }

  int GetNodeIndex(const AnimNodeResource* node_resource) const {
    for (size_t i = 0, n = m_nodes.size(); i < n; i++) {
      if (m_nodes[i] == node_resource) {
        return i;
      }
    }

    std::cerr << "Error: could not find node index for node resource "
              << node_resource << std::endl;
    return -1;
  }

  [[maybe_unused]] size_t AddNode(AnimNodeResource* node_resource) {
    m_nodes.push_back(node_resource);
    m_node_input_connection_indices.emplace_back();
    m_node_inputs_subtree.emplace_back();
    return m_nodes.size() - 1;
  }

  [[nodiscard]] size_t GetNumNodes() const { return m_nodes.size(); }
  [[nodiscard]] AnimNodeResource* GetNode(size_t i) { return m_nodes[i]; }
  [[nodiscard]] const AnimNodeResource* GetNode(size_t i) const {
    return m_nodes[i];
  }
  [[nodiscard]] const std::vector<AnimNodeResource*>& GetNodes() const {
    return m_nodes;
  }

  [[nodiscard]] size_t GetNumConnections() const {
    return m_connections.size();
  }

  [[nodiscard]] BlendTreeConnectionResource* GetConnection(size_t i) {
    return &m_connections[i];
  }

  [[nodiscard]] const BlendTreeConnectionResource* GetConnection(
      size_t i) const {
    return &m_connections[i];
  }

  [[nodiscard]] const std::vector<BlendTreeConnectionResource>& GetConnections()
      const {
    return m_connections;
  }

  Socket* GetNodeOutputSocket(
      const AnimNodeResource* node,
      const std::string& output_socket_name) const;

  const Socket* GetNodeOutputSocketByIndex(
      const AnimNodeResource* node,
      const size_t socket_output_index) const;

  Socket* GetNodeInputSocket(
      const AnimNodeResource* node,
      const std::string& input_socket_name) const;

  const Socket* GetNodeInputSocketByIndex(
      const AnimNodeResource* node,
      const size_t socket_input_index) const;

  std::vector<Socket> GetNodeOutputSockets(const AnimNodeResource* node) const;
  std::vector<Socket> GetNodeInputSockets(const AnimNodeResource* node) const;

  bool ConnectSockets(
      const AnimNodeResource* source_node,
      const std::string& source_socket_name,
      const AnimNodeResource* target_node,
      const std::string& target_socket_name);

  bool DisconnectSockets(
      const AnimNodeResource* source_node,
      const std::string& source_socket_name,
      const AnimNodeResource* target_node,
      const std::string& target_socket_name);

  bool IsConnectionValid(
      const AnimNodeResource* source_node,
      const std::string& source_socket_name,
      const AnimNodeResource* target_node,
      const std::string& target_socket_name) const;

  const BlendTreeConnectionResource* FindConnectionForSocket(
      const AnimNodeResource* node,
      const std::string& socket_name) const;

  bool IsSocketConnected(
      const AnimNodeResource* node,
      const std::string& socket_name) const {
    const BlendTreeConnectionResource* connection =
        FindConnectionForSocket(node, socket_name);
    return connection != nullptr;
  }

  std::vector<Socket*> GetConstantNodeInputs(
      std::vector<NodeDescriptorBase*>& instance_node_descriptors) const {
    std::vector<Socket*> result;

    for (size_t i = 0; i < m_nodes.size(); i++) {
      for (size_t j = 0,
                  num_inputs = instance_node_descriptors[i]->m_inputs.size();
           j < num_inputs;
           j++) {
        Socket& input = instance_node_descriptors[i]->m_inputs[j];

        if (*input.m_reference.ptr_ptr == nullptr) {
          memcpy(
              &input.m_value,
              &m_nodes[i]->m_virtual_socket_accessor->m_inputs[j].m_value,
              sizeof(Socket::SocketValue));
          result.push_back(&input);
        }
      }
    }

    return result;
  }

  size_t GetNodeIndexForOutputSocket(const std::string& socket_name) const {
    for (size_t i = 0; i < m_connections.size(); i++) {
      const BlendTreeConnectionResource& connection = m_connections[i];
      if (connection.target_node_index == 0
          && connection.target_socket_name == socket_name) {
        return connection.source_node_index;
      }
    }

    std::cerr << "Error: could not find a node connected to output '"
              << socket_name << "'." << std::endl;

    return -1;
  }

  size_t GetNodeIndexForInputSocket(const std::string& socket_name) const {
    for (size_t i = 0; i < m_connections.size(); i++) {
      const BlendTreeConnectionResource& connection = m_connections[i];
      if (connection.source_node_index == 1
          && connection.source_socket_name == socket_name) {
        return connection.target_node_index;
      }
    }

    std::cerr << "Error: could not find a node connected to input '"
              << socket_name << "'." << std::endl;

    return -1;
  }

  void UpdateTreeTopologyInfo();

  [[nodiscard]] const std::vector<size_t>& GetNodeEvalOrder() const {
    return m_node_eval_order;
  }

 private:
  void UpdateNodeEvalOrder() {
    m_node_eval_order.clear();
    UpdateNodeEvalOrderRecursive(0);
  }
  void UpdateNodeEvalOrderRecursive(size_t node_index);
  void UpdateNodeSubtrees();

  std::vector<AnimNodeResource*> m_nodes;
  std::vector<BlendTreeConnectionResource> m_connections;
  std::vector<size_t> m_node_eval_order;
};

struct StateMachineTransitionResources {
  size_t source_state_index = -1;
  size_t target_state_index = -1;
  float blend_time = 0.f;
  bool sync_blend = false;
};

struct StateMachineResource {
  std::vector<AnimNodeResource> m_states;
  std::vector<StateMachineTransitionResources> m_transitions;
};

struct AnimGraphResource : AnimNodeResource {
  explicit AnimGraphResource(AnimGraphType graph_type);
  virtual ~AnimGraphResource() { Clear(); };

  std::string m_graph_type_name;

  BlendTreeResource m_blend_tree_resource;
  typedef std::pair<const AnimNodeResource*, std::string> NodeSocketPair;
  typedef std::map<NodeSocketPair, int> NodeSocketDataOffsetMap;

  StateMachineResource m_state_machine_resource;

  void Clear() { m_blend_tree_resource.Reset(); }
  bool SaveToFile(const char* filename) const;
  bool LoadFromFile(const char* filename);

  void CreateBlendTreeInstance(AnimGraphBlendTree& result) const;

  template <typename T>
  bool RegisterBlendTreeInputSocket(const std::string& socket_name) {
    Socket socket;
    socket.m_name = socket_name;
    socket.m_type = GetSocketType<T>();
    socket.m_type_size = sizeof(T);

    return RegisterBlendTreeInputSocket(socket);
  }

  bool RegisterBlendTreeInputSocket(const Socket& socket) {
    AnimNodeResource* input_node = m_blend_tree_resource.GetGraphInputNode();

    Socket* input_socket =
        m_blend_tree_resource.GetGraphInputSocket(socket.m_name.c_str());

    if (input_socket != nullptr) {
      std::cerr << "Error: cannot register output socket as socket with name '"
                << socket.m_name << "' already exists!" << std::endl;
      return false;
    }

    input_node->m_virtual_socket_accessor->m_outputs.push_back(socket);
    m_virtual_socket_accessor->m_inputs =
        input_node->m_virtual_socket_accessor->m_outputs;

    return true;
  }

  template <typename T>
  bool RegisterBlendTreeOutputSocket(const std::string& socket_name) {
    Socket socket;
    socket.m_name = socket_name;
    socket.m_type = GetSocketType<T>();
    socket.m_type_size = sizeof(T);

    return RegisterBlendTreeOutputSocket(socket);
  }

  bool RegisterBlendTreeOutputSocket(const Socket& socket) {
    AnimNodeResource* output_node = m_blend_tree_resource.GetGraphOutputNode();

    Socket* output_socket =
        m_blend_tree_resource.GetGraphOutputSocket(socket.m_name.c_str());

    if (output_socket != nullptr) {
      std::cerr << "Error: cannot register output socket as socket with name '"
                << socket.m_name << "' already exists!" << std::endl;
      return false;
    }

    output_node->m_virtual_socket_accessor->m_inputs.push_back(socket);
    m_virtual_socket_accessor->m_outputs =
        output_node->m_virtual_socket_accessor->m_inputs;

    return true;
  }

  void CreateStateMachineInstance(AnimGraphStateMachine& result) const;

 private:
  // BlendTree
  bool SaveBlendTreeResourceToFile(const char* filename) const;
  void CreateBlendTreeRuntimeNodeInstances(AnimGraphBlendTree& result) const;
  void PrepareBlendTreeIOData(
      AnimGraphBlendTree& instance,
      NodeSocketDataOffsetMap& node_offset_map) const;
  void CreateBlendTreeConnectionInstances(
      AnimGraphBlendTree& instance,
      NodeSocketDataOffsetMap& node_offset_map) const;
  void SetRuntimeNodeProperties(AnimGraphBlendTree& result) const;

  bool SaveStateMachineResourceToFile(const char* filename) const;
  bool LoadStateMachineResourceFromJson(nlohmann::json const& json_data);
};

static inline AnimNodeResource* AnimNodeResourceFactory(
    const std::string& node_type_name) {
  AnimNodeResource* result;

  if (node_type_name == "BlendTree") {
    AnimGraphResource* blend_tree_resource =
        new AnimGraphResource(AnimGraphType::GraphTypeBlendTree);
    result = blend_tree_resource;
  } else {
    result = new AnimNodeResource();
    result->m_virtual_socket_accessor =
        VirtualAnimNodeDescriptorFactory(node_type_name);
  }

  result->m_node_type_name = node_type_name;

  return result;
}

#endif  //ANIMTESTBED_ANIMGRAPHRESOURCE_H