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

#include "AnimGraph.h"

#include <cstring>

bool AnimGraph::init(AnimGraphContext& context) {
  context.m_graph = this;

  for (size_t i = 2; i < m_nodes.size(); i++) {
    if (!m_nodes[i]->Init(context)) {
      return false;
    }
  }

  std::vector<AnimGraphConnection>& graph_outputs = m_node_input_connections[0];
  for (size_t i = 0, n = graph_outputs.size(); i < n; i++) {
    AnimGraphConnection& connection = graph_outputs[i];
    if (connection.m_target_socket.m_type == SocketType::SocketTypeAnimation) {
      AnimData* graph_anim_output =
          static_cast<AnimData*>(connection.m_target_socket.m_reference.ptr);
      assert(graph_anim_output != nullptr);
      graph_anim_output->m_local_matrices.resize(
          context.m_skeleton->num_soa_joints());
    }
  }

  return true;
}

void AnimGraph::updateOrderedNodes() {
  m_eval_ordered_nodes.clear();
  updateOrderedNodesRecursive(0);
}

void AnimGraph::updateOrderedNodesRecursive(int node_index) {
  AnimNode* node = m_nodes[node_index];
  const std::vector<AnimGraphConnection> node_input_connections =
      m_node_input_connections[node_index];
  for (size_t i = 0, n = node_input_connections.size(); i < n; i++) {
    int input_node_index =
        getAnimNodeIndex(node_input_connections[i].m_source_node);

    if (input_node_index == 1) {
      continue;
    }

    updateOrderedNodesRecursive(input_node_index);
  }

  if (node_index != 0) {
    m_eval_ordered_nodes.push_back(node);
  }
}

void AnimGraph::markActiveNodes() {
  for (size_t i = 0, n = m_nodes.size(); i < n; i++) {
    m_nodes[i]->m_state = AnimNodeEvalState::Deactivated;
  }

  const std::vector<AnimGraphConnection>& graph_output_inputs =
      m_node_input_connections[0];
  for (size_t i = 0, n = graph_output_inputs.size(); i < n; i++) {
    const AnimGraphConnection& graph_input = graph_output_inputs[i];
    AnimNode* node = graph_input.m_source_node;
    if (node != nullptr) {
      node->m_state = AnimNodeEvalState::Activated;
    }
  }

  for (size_t i = m_eval_ordered_nodes.size() - 1; i > 0; i--) {
    AnimNode* node = m_eval_ordered_nodes[i];
    if (checkIsNodeActive(node)) {
      int node_index = node->m_index;
      node->MarkActiveInputs(m_node_input_connections[node_index]);

      // Non-animation data inputs are always active.
      for (size_t j = 0, nj = m_node_input_connections[node_index].size();
           j < nj;
           j++) {
        const AnimGraphConnection& input =
            m_node_input_connections[node_index][j];
        if (input.m_source_node != nullptr
            && input.m_target_socket.m_type
                   != SocketType::SocketTypeAnimation) {
          input.m_source_node->m_state = AnimNodeEvalState::Activated;
        }
      }
    }
  }
}

void AnimGraph::prepareNodeEval(
    AnimGraphContext& graph_context,
    size_t node_index) {
  for (size_t i = 0, n = m_node_output_connections[node_index].size(); i < n;
       i++) {
    AnimGraphConnection& output_connection =
        m_node_output_connections[node_index][i];
    if (output_connection.m_source_socket.m_type
        != SocketType::SocketTypeAnimation) {
      continue;
    }

    assert (*output_connection.m_source_socket.m_reference.ptr_ptr == nullptr);

    (*output_connection.m_source_socket.m_reference.ptr_ptr) =
        m_anim_data_allocator.allocate(graph_context.m_skeleton);
  }

  for (size_t i = 0, n = m_node_input_connections[node_index].size(); i < n;
       i++) {
    AnimGraphConnection& input_connection =
        m_node_input_connections[node_index][i];
    if (input_connection.m_source_socket.m_type
        != SocketType::SocketTypeAnimation) {
      continue;
    }

    (*input_connection.m_target_socket.m_reference.ptr_ptr) =
        (*input_connection.m_source_socket.m_reference.ptr_ptr);
  }
}

void AnimGraph::finishNodeEval(size_t node_index) {
  for (size_t i = 0, n = m_node_input_connections[node_index].size(); i < n;
       i++) {
    AnimGraphConnection& input_connection =
        m_node_input_connections[node_index][i];
    if (input_connection.m_source_socket.m_type
        != SocketType::SocketTypeAnimation) {
      continue;
    }

    m_anim_data_allocator.free(static_cast<AnimData*>(
        *input_connection.m_source_socket.m_reference.ptr_ptr));
    (*input_connection.m_source_socket.m_reference.ptr_ptr) = nullptr;
  }
}

void AnimGraph::evalInputNode() {
  for (size_t i = 0, n = m_node_output_connections[1].size(); i < n; i++) {
    AnimGraphConnection& graph_input_connection =
        m_node_output_connections[1][i];

    if (graph_input_connection.m_source_socket.m_type
        != SocketType::SocketTypeAnimation) {
      memcpy(
          *graph_input_connection.m_target_socket.m_reference.ptr_ptr,
          graph_input_connection.m_source_socket.m_reference.ptr,
          sizeof(void*));
      printf("bla");
    } else {
      // TODO: how to deal with anim data outputs?
    }
  }
}

void AnimGraph::evalOutputNode() {
  for (size_t i = 0, n = m_node_input_connections[0].size(); i < n; i++) {
    AnimGraphConnection& graph_output_connection =
        m_node_input_connections[0][i];

    if (graph_output_connection.m_source_socket.m_type
        != SocketType::SocketTypeAnimation) {
      memcpy(
          graph_output_connection.m_target_socket.m_reference.ptr,
          graph_output_connection.m_source_socket.m_reference.ptr,
          graph_output_connection.m_target_socket.m_type_size);
    } else {
      AnimData* source_data = static_cast<AnimData*>(
          *graph_output_connection.m_source_socket.m_reference.ptr_ptr);
      AnimData* target_data = static_cast<AnimData*>(
          graph_output_connection.m_target_socket.m_reference.ptr);
      target_data->m_local_matrices = source_data->m_local_matrices;
    }
  }
}

void AnimGraph::evalSyncTracks() {
  for (size_t i = m_eval_ordered_nodes.size() - 1; i >= 0; i--) {
    AnimNode* node = m_eval_ordered_nodes[i];
    int node_index = node->m_index;
    if (node->m_state == AnimNodeEvalState::Deactivated) {
      continue;
    }

    node->CalcSyncTrack(m_node_input_connections[node_index]);
  }
}

void AnimGraph::updateTime(float dt) {
  const std::vector<AnimGraphConnection>& graph_output_inputs =
      m_node_input_connections[0];
  for (size_t i = 0, n = graph_output_inputs.size(); i < n; i++) {
    AnimNode* node = graph_output_inputs[i].m_source_node;
    if (node != nullptr) {
      node->UpdateTime(node->m_time_now, node->m_time_now + dt);
    }
  }

  for (size_t i = m_eval_ordered_nodes.size() - 1; i > 0; --i) {
    AnimNode* node = m_eval_ordered_nodes[i];
    if (node->m_state != AnimNodeEvalState::TimeUpdated) {
      continue;
    }

    int node_index = node->m_index;
    float node_time_now = node->m_time_now;
    float node_time_last = node->m_time_last;

    const std::vector<AnimGraphConnection>& node_input_connections =
        m_node_input_connections[node_index];
    for (size_t i = 0, n = node_input_connections.size(); i < n; i++) {
      AnimNode* input_node = node_input_connections[i].m_source_node;

      // Only propagate time updates via animation sockets.
      if (input_node != nullptr
          && node_input_connections[i].m_target_socket.m_type
                 == SocketType::SocketTypeAnimation
          && input_node->m_state == AnimNodeEvalState::Activated) {
        input_node->UpdateTime(node_time_last, node_time_now);
      }
    }
  }
}

void AnimGraph::evaluate(AnimGraphContext& context) {
  for (int i = 0, n = m_eval_ordered_nodes.size(); i < n; i++) {
    AnimNode* node = m_eval_ordered_nodes[i];

    if (node->m_state == AnimNodeEvalState::Deactivated) {
      continue;
    }

    prepareNodeEval(context, node->m_index);

    node->Evaluate(context);

    finishNodeEval(node->m_index);
  }

  evalOutputNode();
  finishNodeEval(0);
}

Socket* AnimGraph::getInputSocket(const std::string& name) {
  for (size_t i = 0, n = m_node_output_connections[1].size(); i < n; i++) {
    AnimGraphConnection& connection = m_node_output_connections[1][i];
    if (connection.m_source_socket.m_name == name) {
      return &connection.m_source_socket;
    }
  }

  return nullptr;
}

Socket* AnimGraph::getOutputSocket(const std::string& name) {
  for (size_t i = 0, n = m_node_input_connections[0].size(); i < n; i++) {
    AnimGraphConnection& connection = m_node_input_connections[0][i];
    if (connection.m_target_socket.m_name == name) {
      return &connection.m_target_socket;
    }
  }

  return nullptr;
}

const Socket* AnimGraph::getInputSocket(const std::string& name) const {
  for (size_t i = 0, n = m_node_output_connections[1].size(); i < n; i++) {
    const AnimGraphConnection& connection = m_node_output_connections[1][i];
    if (connection.m_source_socket.m_name == name) {
      return &connection.m_source_socket;
    }
  }

  return nullptr;
}

const Socket* AnimGraph::getOutputSocket(const std::string& name) const {
  for (size_t i = 0, n = m_node_input_connections[0].size(); i < n; i++) {
    const AnimGraphConnection& connection = m_node_input_connections[0][i];
    if (connection.m_target_socket.m_name == name) {
      return &connection.m_target_socket;
    }
  }

  return nullptr;
}