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

#ifndef ANIMTESTBED_ANIMGRAPH_H
#define ANIMTESTBED_ANIMGRAPH_H

#include "AnimGraphData.h"
#include "AnimGraphNodes.h"

//
// AnimGraph (Runtime)
//
struct AnimGraph {
  AnimData m_local_transforms;

  std::vector<AnimNode*> m_nodes;
  std::vector<AnimNode*> m_eval_ordered_nodes;
  std::vector<std::vector<AnimGraphConnection> > m_node_input_connections;
  std::vector<std::vector<AnimGraphConnection> > m_node_output_connections;
  std::vector<AnimData*> m_animdata_blocks;
  NodeDescriptorBase* m_node_descriptor;
  char* m_input_buffer = nullptr;
  char* m_output_buffer = nullptr;
  char* m_connection_data_storage = nullptr;

  std::vector<Socket>& getGraphOutputs() { return m_node_descriptor->m_inputs; }
  std::vector<Socket>& getGraphInputs() { return m_node_descriptor->m_outputs; }

  AnimDataAllocator m_anim_data_allocator;

  ~AnimGraph() { dealloc(); }

  bool init(AnimGraphContext& context);
  void dealloc() {
    for (size_t i = 0; i < m_animdata_blocks.size(); i++) {
      m_animdata_blocks[i]->m_local_matrices.vector::~vector();
    }
    m_animdata_blocks.clear();

    m_node_input_connections.clear();
    m_node_output_connections.clear();

    delete[] m_input_buffer;
    delete[] m_output_buffer;
    delete[] m_connection_data_storage;

    for (int i = 0; i < m_nodes.size(); i++) {
      delete m_nodes[i];
    }
    m_nodes.clear();

    delete m_node_descriptor;
  }

  void updateOrderedNodes();
  void updateOrderedNodesRecursive(int node_index);
  void markActiveNodes();
  bool checkIsNodeActive(AnimNode* node) {
    return node->m_state != AnimNodeEvalState::Deactivated;
  }

  void evalSyncTracks();
  void updateTime(float dt);
  void evaluate(AnimGraphContext& context);
  void resetNodeStates() {
    for (size_t i = 0, n = m_nodes.size(); i < n; i++) {
      m_nodes[i]->m_time_now = 0.f;
      m_nodes[i]->m_time_last = 0.f;
      m_nodes[i]->m_state = AnimNodeEvalState::Undefined;
    }
  }

  Socket* getInputSocket(const std::string& name);
  Socket* getOutputSocket(const std::string& name);

  const Socket* getInputSocket(const std::string& name) const;
  const Socket* getOutputSocket(const std::string& name) const;

  /** Sets the address that is used for the specified AnimGraph input Socket.
   *
   * @tparam T Type of the Socket.
   * @param name Name of the Socket.
   * @param value_ptr Pointer where the input is fetched during evaluation.
   */
  template <typename T>
  void SetInput(const char* name, T* value_ptr) {
    m_node_descriptor->SetOutput(name, value_ptr);

    for (int i = 0; i < m_node_output_connections[1].size(); i++) {
      const AnimGraphConnection& graph_input_connection =
          m_node_output_connections[1][i];

      if (graph_input_connection.m_source_socket.m_name == name) {
        *graph_input_connection.m_target_socket.m_reference.ptr_ptr = value_ptr;
      }
    }
  }

  /** Sets the address that is used for the specified AnimGraph output Socket.
   *
   * @tparam T Type of the Socket.
   * @param name Name of the Socket.
   * @param value_ptr Pointer where the graph output output is written to at the end of evaluation.
   */
  template <typename T>
  void SetOutput(const char* name, T* value_ptr) {
    m_node_descriptor->SetInput(name, value_ptr);

    for (int i = 0; i < m_node_input_connections[0].size(); i++) {
      const AnimGraphConnection& graph_output_connection =
          m_node_input_connections[0][i];

      if (graph_output_connection.m_target_socket.m_name == name) {
        if (graph_output_connection.m_source_node == m_nodes[1]
            && graph_output_connection.m_target_node == m_nodes[0]) {
          std::cerr << "Error: cannot set output for direct graph input to graph "
                       "output connections. Use GetOutptPtr for output instead!"
                    << std::endl;

          return;
        }

        *graph_output_connection.m_source_socket.m_reference.ptr_ptr =
            value_ptr;

        // Make sure all other output connections of this pin use the same output pointer
        int source_node_index = getAnimNodeIndex(graph_output_connection.m_source_node);
        for (int j = 0; j < m_node_output_connections[source_node_index].size(); j++) {
          const AnimGraphConnection& source_output_connection = m_node_output_connections[source_node_index][j];
          if (source_output_connection.m_target_node == m_nodes[0]) {
            continue;
          }

          if (source_output_connection.m_source_socket.m_name == graph_output_connection.m_source_socket.m_name) {
            *source_output_connection.m_target_socket.m_reference.ptr_ptr = value_ptr;
          }
        }
      }
    }
  }

  /** Returns the address that is used for the specified AnimGraph output Socket.
   *
   * This function is needed for connections that directly connect an AnimGraph
   * input Socket to an output Socket of the same AnimGraph.
   *
   * @tparam T Type of the Socket.
   * @param name Name of the Socket.
   * @return Address that is used for the specified AnimGraph output Socket.
   */
  template <typename T>
  T* GetOutputPtr(const char* name) {
    for (int i = 0; i < m_node_input_connections[0].size(); i++) {
      const AnimGraphConnection& graph_output_connection =
          m_node_input_connections[0][i];
      if (graph_output_connection.m_target_socket.m_name == name) {
        return static_cast<float*>(*graph_output_connection.m_source_socket.m_reference.ptr_ptr);
      }
    }

    return nullptr;
  }

  void* getInputPtr(const std::string& name) const {
    const Socket* input_socket = getInputSocket(name);
    if (input_socket != nullptr) {
      return input_socket->m_reference.ptr;
    }

    return nullptr;
  }

  void* getOutputPtr(const std::string& name) const {
    const Socket* input_socket = getOutputSocket(name);
    if (input_socket != nullptr) {
      return input_socket->m_reference.ptr;
    }

    return nullptr;
  }

  int getNodeEvalOrderIndex(const AnimNode* node) {
    for (size_t i = 0, n = m_eval_ordered_nodes.size(); i < n; i++) {
      if (m_eval_ordered_nodes[i] == node) {
        return i;
      }
    }

    return -1;
  }
  const AnimNode* getAnimNodeForInput(
      size_t node_index,
      const std::string& input_name) const {
    assert(node_index < m_nodes.size());

    const std::vector<AnimGraphConnection>& input_connection =
        m_node_input_connections[node_index];
    for (size_t i = 0, n = input_connection.size(); i < n; i++) {
      if (input_connection[i].m_target_socket.m_name == input_name) {
        return input_connection[i].m_source_node;
      }
    }

    return nullptr;
  }

  AnimNode* getAnimNode(const char* name) {
    for (size_t i = 0; i < m_nodes.size(); i++) {
      if (m_nodes[i]->m_name == name) {
        return m_nodes[i];
      }
    }

    return nullptr;
  }

  size_t getAnimNodeIndex(AnimNode* node) {
    for (size_t i = 0; i < m_nodes.size(); i++) {
      if (m_nodes[i] == node) {
        return i;
      }
    }

    return -1;
  }
};

#endif  //ANIMTESTBED_ANIMGRAPH_H