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

#include "AnimGraphResource.h"

#include <cstring>
#include <fstream>

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

using json = nlohmann::json;

// forward declarations
static json sAnimGraphResourceBlendTreeToJson(
    const AnimGraphResource& anim_graph_resource);
static bool sAnimGraphResourceBlendTreeFromJson(
    const json& json_data,
    AnimGraphResource* result_graph_resource);

//
// Socket <-> json
//
std::string sSocketTypeToStr(SocketType pin_type) {
  if (pin_type < SocketType::SocketTypeUndefined
      || pin_type >= SocketType::SocketTypeLast) {
    return "Unknown";
  }

  return SocketTypeNames[static_cast<int>(pin_type)];
}

json sSocketToJson(const Socket& socket) {
  json result;
  result["name"] = socket.m_name;
  result["type"] = sSocketTypeToStr(socket.m_type);

  if (socket.m_type == SocketType::SocketTypeString
      && !socket.m_value_string.empty()) {
    result["value"] = socket.m_value_string;
  } else if (socket.m_value.flag) {
    if (socket.m_type == SocketType::SocketTypeBool) {
      result["value"] = socket.m_value.flag;
    } else if (socket.m_type == SocketType::SocketTypeAnimation) {
    } else if (socket.m_type == SocketType::SocketTypeInt) {
      result["value"] = socket.m_value.int_value;
    } else if (socket.m_type == SocketType::SocketTypeFloat) {
      result["value"] = socket.m_value.float_value;
    } else if (socket.m_type == SocketType::SocketTypeVec3) {
      result["value"][0] = socket.m_value.vec3.v[0];
      result["value"][1] = socket.m_value.vec3.v[1];
      result["value"][2] = socket.m_value.vec3.v[2];
    } else if (socket.m_type == SocketType::SocketTypeQuat) {
      result["value"][0] = socket.m_value.quat.v[0];
      result["value"][1] = socket.m_value.quat.v[1];
      result["value"][2] = socket.m_value.quat.v[2];
      result["value"][3] = socket.m_value.quat.v[3];
    } else {
      std::cerr << "Invalid socket type '" << static_cast<int>(socket.m_type)
                << "'." << std::endl;
    }
  }
  return result;
}

Socket sJsonToSocket(const json& json_data) {
  Socket result;
  result.m_type = SocketType::SocketTypeUndefined;
  result.m_reference.ptr = &result.m_value.int_value;
  result.m_name = json_data["name"];

  std::string type_string = json_data["type"];
  bool have_value = json_data.contains("value");

  if (type_string == "Bool") {
    result.m_type = SocketType::SocketTypeBool;
    result.m_type_size = sizeof(bool);
    result.m_reference.ptr = &result.m_value.int_value;

    if (have_value) {
      result.m_value.flag = json_data["value"];
    }
  } else if (type_string == "Animation") {
    result.m_type = SocketType::SocketTypeAnimation;
    result.m_type_size = sizeof(AnimData);
  } else if (type_string == "Int") {
    result.m_type = SocketType::SocketTypeInt;
    result.m_type_size = sizeof(int);
    if (have_value) {
      result.m_value.int_value = json_data["value"];
    }
  } else if (type_string == "Float") {
    result.m_type = SocketType::SocketTypeFloat;
    result.m_type_size = sizeof(float);
    if (have_value) {
      result.m_value.float_value = json_data["value"];
    }
  } else if (type_string == "Vec3") {
    result.m_type = SocketType::SocketTypeVec3;
    result.m_type_size = sizeof(Vec3);
    if (have_value) {
      result.m_value.vec3.x = json_data["value"][0];
      result.m_value.vec3.y = json_data["value"][1];
      result.m_value.vec3.z = json_data["value"][2];
    }
  } else if (type_string == "Quat") {
    result.m_type = SocketType::SocketTypeQuat;
    result.m_type_size = sizeof(Quat);
    if (have_value) {
      result.m_value.quat.x = json_data["value"][0];
      result.m_value.quat.y = json_data["value"][1];
      result.m_value.quat.z = json_data["value"][2];
      result.m_value.quat.w = json_data["value"][3];
    }
  } else if (type_string == "String") {
    result.m_type = SocketType::SocketTypeString;
    result.m_type_size = sizeof(std::string);
    if (have_value) {
      result.m_value_string = json_data["value"];
    }
  } else {
    std::cerr << "Invalid socket type '" << type_string << "'." << std::endl;
  }

  return result;
}

//
// AnimGraphNode <-> json
//
json sAnimGraphNodeToJson(
    const AnimNodeResource* node,
    size_t node_index,
    const std::vector<BlendTreeConnectionResource>& connections) {
  json result;

  result["name"] = node->m_name;
  result["type"] = "AnimNodeResource";
  result["node_type"] = node->m_node_type_name;

  for (size_t j = 0; j < 2; j++) {
    result["position"][j] = node->m_position[j];
  }

  if (node->m_node_type_name == "BlendTree") {
  }

  for (const auto& socket : node->m_socket_accessor->m_inputs) {
    if (socket.m_type == SocketType::SocketTypeAnimation) {
      continue;
    }

    bool socket_connected = false;
    for (const auto& connection : connections) {
      if (connection.source_node_index == node_index
          && connection.source_socket_name == socket.m_name) {
        socket_connected = true;
        break;
      }
    }

    if (!socket_connected) {
      result["inputs"].push_back(sSocketToJson(socket));
    }
  }

  for (auto& property : node->m_socket_accessor->m_properties) {
    result["properties"][property.m_name] = sSocketToJson(property);
  }

  return result;
}

AnimNodeResource* sAnimGraphNodeFromJson(
    const json& json_node,
    size_t node_index) {
  std::string node_type = json_node["node_type"];

  if (node_type == "BlendTree") {
    AnimGraphResource* result = new AnimGraphResource();
    sAnimGraphResourceBlendTreeFromJson(json_node, result);
    return result;
  }

  AnimNodeResource* result;
  result = new AnimNodeResource();
  result->m_name = json_node["name"];
  result->m_node_type_name = node_type;
  result->m_position[0] = json_node["position"][0];
  result->m_position[1] = json_node["position"][1];

  result->m_anim_node = AnimNodeFactory(result->m_node_type_name);
  result->m_socket_accessor =
      AnimNodeDescriptorFactory(result->m_node_type_name, result->m_anim_node);

  for (auto& property : result->m_socket_accessor->m_properties) {
    property = sJsonToSocket(json_node["properties"][property.m_name]);
  }

  if (node_index != 0 && node_index != 1 && json_node.contains("inputs")) {
    for (size_t j = 0, n = json_node["inputs"].size(); j < n; j++) {
      assert(json_node["inputs"][j].contains("name"));
      std::string input_name = json_node["inputs"][j]["name"];
      Socket* input_socket =
          result->m_socket_accessor->GetInputSocket(input_name.c_str());
      if (input_socket == nullptr) {
        std::cerr << "Could not find input socket with name " << input_name
                  << " for node type " << result->m_node_type_name << std::endl;
        abort();
      }
      *input_socket = sJsonToSocket(json_node["inputs"][j]);
    }
  }

  return result;
}

//
// AnimGraphConnectionResource <-> Json
//
json sAnimGraphConnectionToJson(const BlendTreeConnectionResource& connection) {
  json result;

  result["type"] = "AnimGraphConnectionResource";

  result["source_node_index"] = connection.source_node_index;
  result["source_socket_name"] = connection.source_socket_name;

  result["target_node_index"] = connection.target_node_index;
  result["target_socket_name"] = connection.target_socket_name;

  return result;
}

BlendTreeConnectionResource sAnimGraphConnectionFromJson(
    const json& json_node) {
  BlendTreeConnectionResource connection;

  connection.source_node_index = json_node["source_node_index"];
  connection.source_socket_name = json_node["source_socket_name"];

  connection.target_node_index = json_node["target_node_index"];
  connection.target_socket_name = json_node["target_socket_name"];

  return connection;
}

static json sAnimGraphResourceBlendTreeToJson(
    const AnimGraphResource& anim_graph_resource) {
  json result;

  result["name"] = anim_graph_resource.m_name;
  result["type"] = "AnimNodeResource";
  result["node_type"] = "BlendTree";

  const BlendTreeResource& blend_tree_resource =
      anim_graph_resource.m_blend_tree_resource;

  for (size_t i = 0; i < blend_tree_resource.m_nodes.size(); i++) {
    const AnimNodeResource* node = blend_tree_resource.m_nodes[i];

    if (node->m_node_type_name == "BlendTree") {
      const AnimGraphResource* graph_resource =
          dynamic_cast<const AnimGraphResource*>(node);
      result["nodes"][i] = sAnimGraphResourceBlendTreeToJson(*graph_resource);
    } else {
      result["nodes"][i] =
          sAnimGraphNodeToJson(node, i, blend_tree_resource.m_connections);
    }
  }

  for (size_t i = 0; i < blend_tree_resource.m_connections.size(); i++) {
    const BlendTreeConnectionResource& connection =
        blend_tree_resource.m_connections[i];
    result["connections"][i] = sAnimGraphConnectionToJson(connection);
  }

  // Graph inputs and outputs
  {
    const AnimNodeResource* graph_output_node = blend_tree_resource.m_nodes[0];
    const std::vector<Socket> graph_inputs =
        graph_output_node->m_socket_accessor->m_inputs;
    for (size_t i = 0; i < graph_inputs.size(); i++) {
      result["nodes"][0]["inputs"][i] = sSocketToJson(graph_inputs[i]);
    }

    const AnimNodeResource* graph_input_node = blend_tree_resource.m_nodes[1];
    const std::vector<Socket> graph_outputs =
        graph_input_node->m_socket_accessor->m_outputs;
    for (size_t i = 0; i < graph_outputs.size(); i++) {
      result["nodes"][1]["outputs"][i] = sSocketToJson(graph_outputs[i]);
    }
  }
  return result;
}

static bool sAnimGraphResourceBlendTreeFromJson(
    const json& json_data,
    AnimGraphResource* result_graph_resource) {
  BlendTreeResource& blend_tree_resource =
      result_graph_resource->m_blend_tree_resource;

  result_graph_resource->m_graph_type_name = "BlendTree";
  result_graph_resource->m_node_type_name = "BlendTree";
  result_graph_resource->m_name = json_data["name"];

  // Load nodes
  for (size_t i = 0, n = json_data["nodes"].size(); i < n; i++) {
    const json& json_node = json_data["nodes"][i];
    if (json_node["type"] != "AnimNodeResource") {
      std::cerr
          << "Invalid json object. Expected type 'AnimNodeResource' but got '"
          << json_node["type"] << "'." << std::endl;
      return false;
    }

    AnimNodeResource* node = sAnimGraphNodeFromJson(json_node, i);
    blend_tree_resource.m_nodes.push_back(node);
  }

  // Graph outputs
  const json& graph_outputs = json_data["nodes"][0]["inputs"];
  for (const auto& graph_output : graph_outputs) {
    AnimNodeResource* graph_node = blend_tree_resource.m_nodes[0];
    graph_node->m_socket_accessor->m_inputs.push_back(
        sJsonToSocket(graph_output));
  }

  // Graph inputs (optional)
  if (json_data["nodes"][1].contains("outputs")) {
    const json& graph_inputs = json_data["nodes"][1]["outputs"];
    for (const auto& graph_input : graph_inputs) {
      AnimNodeResource* graph_node = blend_tree_resource.m_nodes[1];
      graph_node->m_socket_accessor->m_outputs.push_back(
          sJsonToSocket(graph_input));
    }
  }

  // Load connections
  for (const auto& json_connection : json_data["connections"]) {
    if (json_connection["type"] != "AnimGraphConnectionResource") {
      std::cerr << "Invalid json object. Expected type "
                   "'AnimGraphConnectionResource' "
                   "but got '"
                << json_connection["type"] << "'." << std::endl;
      return false;
    }

    BlendTreeConnectionResource connection =
        sAnimGraphConnectionFromJson(json_connection);
    blend_tree_resource.m_connections.push_back(connection);
  }

  return true;
}

bool BlendTreeResource::ConnectSockets(
    const AnimNodeResource* source_node,
    const std::string& source_socket_name,
    const AnimNodeResource* target_node,
    const std::string& target_socket_name) {
  size_t source_node_index = GetNodeIndex(source_node);
  size_t target_node_index = GetNodeIndex(target_node);

  if (source_node_index >= m_nodes.size()
      || target_node_index >= m_nodes.size()) {
    std::cerr << "Cannot connect nodes: could not find nodes." << std::endl;
    return false;
  }

  Socket* source_socket;
  Socket* target_socket;

  if (target_node->m_node_type_name == "BlendTree") {
    const AnimGraphResource* target_graph_resource =
        dynamic_cast<const AnimGraphResource*>(target_node);
    AnimNodeResource* graph_output_node =
        target_graph_resource->m_blend_tree_resource.GetGraphInputNode();
    target_socket = graph_output_node->m_socket_accessor->GetOutputSocket(
        target_socket_name.c_str());
  } else {
    target_socket = target_node->m_socket_accessor->GetInputSocket(
        target_socket_name.c_str());
  }

  if (source_node->m_node_type_name == "BlendTree") {
    const AnimGraphResource* source_graph_resource =
        dynamic_cast<const AnimGraphResource*>(source_node);
    AnimNodeResource* graph_output_node =
        source_graph_resource->m_blend_tree_resource.GetGraphOutputNode();
    source_socket = graph_output_node->m_socket_accessor->GetInputSocket(
        source_socket_name.c_str());
  } else {
    source_socket = source_node->m_socket_accessor->GetOutputSocket(
        source_socket_name.c_str());
  }

  if (source_socket == nullptr) {
    std::cerr << "Cannot connect nodes: could not find source socket '"
              << source_socket_name << "'." << std::endl;
  }

  if (target_socket == nullptr) {
    std::cerr << "Cannot connect nodes: could not find target socket '"
              << target_socket_name << "'." << std::endl;
  }

  if (target_socket == nullptr || source_socket == nullptr) {
    return false;
  }

  BlendTreeConnectionResource connection;
  connection.source_node_index = source_node_index;
  connection.source_socket_name = source_socket_name;
  connection.target_node_index = target_node_index;
  connection.target_socket_name = target_socket_name;
  m_connections.push_back(connection);

  return true;
}

bool AnimGraphResource::LoadFromFile(const char* filename) {
  std::ifstream input_file;
  input_file.open(filename);
  std::stringstream buffer;
  buffer << input_file.rdbuf();

  json json_data = json::parse(buffer.str(), nullptr, false);
  if (json_data.is_discarded()) {
    std::cerr << "Error parsing json of file '" << filename << "'."
              << std::endl;

    return false;
  }

  if (json_data["type"] != "AnimNodeResource") {
    std::cerr
        << "Invalid json object. Expected type 'AnimNodeResource' but got '"
        << json_data["type"] << "'." << std::endl;

    return false;
  }

  if (json_data["node_type"] == "BlendTree") {
    return sAnimGraphResourceBlendTreeFromJson(json_data, this);
  } else if (json_data["node_type"] == "StateMachine") {
    return LoadStateMachineResourceFromJson(json_data);
  }

  std::cerr << "Invalid node_type. Expected type 'BlendTree' or "
               "'StateMachine' but got '"
            << json_data["node_type"] << "'." << std::endl;

  return false;
}

bool AnimGraphResource::SaveToFile(const char* filename) const {
  if (m_graph_type_name == "BlendTree") {
    return SaveBlendTreeResourceToFile(filename);
  } else if (m_graph_type_name == "StateMachine") {
    return SaveStateMachineResourceToFile(filename);
  }

  std::cerr << "Invalid AnimGraphResource type: " << m_graph_type_name << "."
            << std::endl;

  return false;
}

bool AnimGraphResource::SaveBlendTreeResourceToFile(
    const char* filename) const {
  json result;

  result = sAnimGraphResourceBlendTreeToJson(*this);

  std::ofstream output_file;
  output_file.open(filename);
  output_file << result.dump(4, ' ') << std::endl;
  output_file.close();

  return true;
}

void AnimGraphResource::CreateBlendTreeInstance(
    AnimGraphBlendTree& result) const {
  if (m_node_type_name != "BlendTree") {
    std::cerr
        << "Invalid AnimGraphResource. Expected type 'BlendTree' but got '"
        << m_graph_type_name << "'." << std::endl;
    return;
  }

  result.m_name = m_name;

  CreateBlendTreeRuntimeNodeInstances(result);
  PrepareBlendTreeIOData(result);
  SetRuntimeNodeProperties(result);
  CreateBlendTreeConnectionInstances(result);

  result.UpdateOrderedNodes();
  result.ResetNodeStates();
}

void AnimGraphResource::CreateBlendTreeRuntimeNodeInstances(
    AnimGraphBlendTree& result) const {
  for (auto node_resource : m_blend_tree_resource.m_nodes) {
    AnimNode* node = AnimNodeFactory(node_resource->m_node_type_name);

    if (node_resource->m_node_type_name == "BlendTree") {
      AnimGraphResource* embedded_blend_tree_resource =
          dynamic_cast<AnimGraphResource*>(node_resource);
      assert(embedded_blend_tree_resource != nullptr);
      AnimGraphBlendTree* embedded_blend_tree =
          dynamic_cast<AnimGraphBlendTree*>(node);
      assert(embedded_blend_tree != nullptr);

      embedded_blend_tree_resource->CreateBlendTreeInstance(
          *embedded_blend_tree);
      embedded_blend_tree_resource->m_socket_accessor->m_inputs =
          embedded_blend_tree->m_node_descriptor->m_outputs;
      embedded_blend_tree_resource->m_socket_accessor->m_outputs =
          embedded_blend_tree->m_node_descriptor->m_inputs;
    }

    node->m_name = node_resource->m_name;
    node->m_node_type_name = node_resource->m_node_type_name;
    result.m_nodes.push_back(node);

    // runtime node connections
    result.m_node_input_connections.emplace_back();
    result.m_node_output_connections.emplace_back();
  }
}

void AnimGraphResource::PrepareBlendTreeIOData(
    AnimGraphBlendTree& instance) const {
  instance.m_node_descriptor =
      AnimNodeDescriptorFactory("BlendTree", instance.m_nodes[0]);

  instance.m_node_descriptor->m_outputs =
      m_blend_tree_resource.m_nodes[1]->m_socket_accessor->m_outputs;
  instance.m_node_descriptor->m_inputs =
      m_blend_tree_resource.m_nodes[0]->m_socket_accessor->m_inputs;

  //
  // graph inputs
  //
  int input_block_size = 0;
  std::vector<Socket>& graph_inputs = instance.GetGraphInputs();
  for (int i = 0; i < graph_inputs.size(); i++) {
    input_block_size += sizeof(void*);
  }

  if (input_block_size > 0) {
    instance.m_input_buffer = new char[input_block_size];
    memset(instance.m_input_buffer, 0, input_block_size);
  }

  int input_block_offset = 0;
  for (int i = 0; i < graph_inputs.size(); i++) {
    graph_inputs[i].m_reference.ptr =
        (void*)&instance.m_input_buffer[input_block_offset];
    instance.m_node_descriptor->m_outputs[i].m_reference.ptr =
        &instance.m_input_buffer[input_block_offset];
    input_block_offset += sizeof(void*);
  }

  //
  // graph outputs
  //
  int output_block_size = 0;
  std::vector<Socket>& graph_outputs = instance.GetGraphOutputs();
  for (int i = 0; i < graph_outputs.size(); i++) {
    output_block_size += sizeof(void*);
  }

  if (output_block_size > 0) {
    instance.m_output_buffer = new char[output_block_size];
    memset(instance.m_output_buffer, 0, output_block_size);
  }

  int output_block_offset = 0;
  for (int i = 0; i < graph_outputs.size(); i++) {
    instance.m_node_descriptor->m_inputs[i].m_reference.ptr =
        &instance.m_output_buffer[output_block_offset];
    output_block_offset += sizeof(void*);
  }

  // connections: make source and target sockets point to the same address in the connection data storage.
  // TODO: instead of every connection, only create data blocks for the source sockets and make sure every source socket gets allocated once.
  size_t connection_data_storage_size = 0;
  for (const auto& connection : m_blend_tree_resource.m_connections) {
    const AnimNodeResource* source_node =
        m_blend_tree_resource.m_nodes[connection.source_node_index];
    Socket* source_socket = source_node->m_socket_accessor->GetOutputSocket(
        connection.source_socket_name.c_str());
    connection_data_storage_size += source_socket->m_type_size;
  }

  if (connection_data_storage_size > 0) {
    instance.m_connection_data_storage = new char[connection_data_storage_size];
    memset(instance.m_connection_data_storage, 0, connection_data_storage_size);
  }
}

void AnimGraphResource::CreateBlendTreeConnectionInstances(
    AnimGraphBlendTree& instance) const {
  std::vector<NodeDescriptorBase*> instance_node_descriptors(
      m_blend_tree_resource.m_nodes.size(),
      nullptr);
  for (int i = 0; i < m_blend_tree_resource.m_nodes.size(); i++) {
    instance_node_descriptors[i] = AnimNodeDescriptorFactory(
        m_blend_tree_resource.m_nodes[i]->m_node_type_name,
        instance.m_nodes[i]);

    if (i > 1
        && m_blend_tree_resource.m_nodes[i]->m_node_type_name == "BlendTree") {
      instance_node_descriptors[i]->m_inputs =
          m_blend_tree_resource.m_nodes[i]->m_socket_accessor->m_inputs;
      instance_node_descriptors[i]->m_outputs =
          m_blend_tree_resource.m_nodes[i]->m_socket_accessor->m_outputs;
    }
  }

  instance_node_descriptors[0]->m_inputs = instance.m_node_descriptor->m_inputs;
  instance_node_descriptors[1]->m_outputs =
      instance.m_node_descriptor->m_outputs;

  size_t connection_data_offset = 0;
  for (const auto& connection : m_blend_tree_resource.m_connections) {
    NodeDescriptorBase* source_node_descriptor =
        instance_node_descriptors[connection.source_node_index];
    NodeDescriptorBase* target_node_descriptor =
        instance_node_descriptors[connection.target_node_index];

    AnimNode* source_node = instance.m_nodes[connection.source_node_index];
    AnimNode* target_node = instance.m_nodes[connection.target_node_index];

    Socket* source_socket = source_node_descriptor->GetOutputSocket(
        connection.source_socket_name.c_str());
    Socket* target_socket = target_node_descriptor->GetInputSocket(
        connection.target_socket_name.c_str());

    AnimGraphConnection instance_connection;

    if (source_node->m_node_type_name == "BlendTree") {
      // For embedded subgraphs we have to ensure two things:
      // 1. The parent graph target node must activate the source node within
      //    the embedded subgraph.
      // 2. The parent graph target node must activate the AnimNode of the
      //    parent graph which contains the embedded subgraph. Otherwise, the
      //    embedded graph does not get evaluated.
      // For each case we have to add a connection here. First we insert the
      // default connection within the graph (i.e. for case 2). Then we alter
      // the source node such that it points to the node within the embedded
      // subgraph.
      AnimGraphConnection embedded_graph_activation_connection;
      embedded_graph_activation_connection.m_source_node = source_node;
      embedded_graph_activation_connection.m_source_socket = *source_socket;
      embedded_graph_activation_connection.m_target_node = target_node;
      embedded_graph_activation_connection.m_target_socket = *target_socket;
      instance.m_node_input_connections[connection.target_node_index].push_back(
          embedded_graph_activation_connection);
      instance.m_node_output_connections[connection.source_node_index]
          .push_back(embedded_graph_activation_connection);

      AnimGraphResource* source_blend_tree_resource =
          dynamic_cast<AnimGraphResource*>(
              m_blend_tree_resource.m_nodes[connection.source_node_index]);
      AnimGraphBlendTree* source_blend_tree =
          dynamic_cast<AnimGraphBlendTree*>(source_node);

      size_t source_blend_tree_output_node_index =
          source_blend_tree_resource->m_blend_tree_resource
              .GetNodeIndexForOutputSocket(connection.source_socket_name);
      source_node =
          source_blend_tree->m_nodes[source_blend_tree_output_node_index];
      instance_connection.m_crosses_hierarchy = true;
    } else if (target_node->m_node_type_name == "BlendTree") {
      // When a connection points to an embedded blend tree we have to ensure
      // that the embedded node knows about its connection partner in the parent
      // tree. This allows the embedded node to properly activate the node in
      // the parent graph.
      AnimGraphResource* target_blend_tree_resource =
          dynamic_cast<AnimGraphResource*>(
              m_blend_tree_resource.m_nodes[connection.target_node_index]);
      AnimGraphBlendTree* target_blend_tree =
          dynamic_cast<AnimGraphBlendTree*>(target_node);

      size_t target_blend_tree_output_node_index =
          target_blend_tree_resource->m_blend_tree_resource
              .GetNodeIndexForInputSocket(connection.target_socket_name);
      target_node =
          target_blend_tree->m_nodes[target_blend_tree_output_node_index];

      // Make sure that the embedded node input activates the parent graph node.
      std::vector<AnimGraphConnection>& embeded_target_node_inputs =
          target_blend_tree
              ->m_node_input_connections[target_blend_tree_output_node_index];
      for (size_t j = 0; j < embeded_target_node_inputs.size(); j++) {
        AnimGraphConnection& embedded_target_connection =
            embeded_target_node_inputs[j];
        if (embedded_target_connection.m_source_socket.m_name
            == target_socket->m_name) {
          embedded_target_connection.m_source_node = source_node;
          embedded_target_connection.m_crosses_hierarchy = true;

          // In addition: make sure we expose the embedded socket to the
          // connection in the parent blend tree. That way
          // parent_tree.SetValue<>() correctly propagates the value to the
          // embedded node socket.
          target_socket = &embedded_target_connection.m_target_socket;
        }
      }
    }

    instance_connection.m_source_node = source_node;
    instance_connection.m_source_socket = *source_socket;
    instance_connection.m_target_node = target_node;
    instance_connection.m_target_socket = *target_socket;
    instance.m_node_input_connections[connection.target_node_index].push_back(
        instance_connection);
    instance.m_node_output_connections[connection.source_node_index].push_back(
        instance_connection);

    source_node_descriptor->SetOutputUnchecked(
        connection.source_socket_name.c_str(),
        &instance.m_connection_data_storage[connection_data_offset]);

    target_node_descriptor->SetInputUnchecked(
        connection.target_socket_name.c_str(),
        &instance.m_connection_data_storage[connection_data_offset]);

    if (source_socket->m_type == SocketType::SocketTypeAnimation) {
      instance.m_animdata_blocks.push_back(
          (AnimData*)(&instance
                           .m_connection_data_storage[connection_data_offset]));
    }

    connection_data_offset += source_socket->m_type_size;
  }

  //
  // const node inputs
  //
  std::vector<Socket*> const_inputs =
      m_blend_tree_resource.GetConstantNodeInputs(instance_node_descriptors);
  size_t const_node_inputs_buffer_size = 0;
  for (auto& const_input : const_inputs) {
    if (const_input->m_type == SocketType::SocketTypeString) {
      // TODO: implement string const node input support
      std::cerr << "Error: const inputs for strings not yet implemented!"
                << std::endl;
      abort();
    }
    const_node_inputs_buffer_size += const_input->m_type_size;
  }

  if (const_node_inputs_buffer_size > 0) {
    instance.m_const_node_inputs = new char[const_node_inputs_buffer_size];
    memset(instance.m_const_node_inputs, '\0', const_node_inputs_buffer_size);
  }

  size_t const_input_buffer_offset = 0;
  for (auto& i : const_inputs) {
    Socket* const_input = i;

    // TODO: implement string const node input support
    assert(const_input->m_type != SocketType::SocketTypeString);

    *const_input->m_reference.ptr_ptr =
        &instance.m_const_node_inputs[const_input_buffer_offset];
    memcpy(
        *const_input->m_reference.ptr_ptr,
        &const_input->m_value,
        i->m_type_size);

    const_input_buffer_offset += i->m_type_size;
  }

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

void AnimGraphResource::SetRuntimeNodeProperties(
    AnimGraphBlendTree& result) const {
  for (int i = 2; i < m_blend_tree_resource.m_nodes.size(); i++) {
    const AnimNodeResource* node_resource = m_blend_tree_resource.m_nodes[i];

    NodeDescriptorBase* node_instance_accessor = AnimNodeDescriptorFactory(
        node_resource->m_node_type_name,
        result.m_nodes[i]);

    std::vector<Socket>& resource_properties =
        node_resource->m_socket_accessor->m_properties;
    for (const auto& property : resource_properties) {
      const std::string& name = property.m_name;

      switch (property.m_type) {
        case SocketType::SocketTypeBool:
          node_instance_accessor->SetProperty(
              name.c_str(),
              property.m_value.flag);
          break;
        case SocketType::SocketTypeInt:
          node_instance_accessor->SetProperty(
              name.c_str(),
              property.m_value.int_value);
          break;
        case SocketType::SocketTypeFloat:
          node_instance_accessor->SetProperty(
              name.c_str(),
              property.m_value.float_value);
          break;
        case SocketType::SocketTypeVec3:
          node_instance_accessor->SetProperty<Vec3>(
              name.c_str(),
              property.m_value.vec3);
          break;
        case SocketType::SocketTypeQuat:
          node_instance_accessor->SetProperty(
              name.c_str(),
              property.m_value.quat);
          break;
        case SocketType::SocketTypeString:
          node_instance_accessor->SetProperty(
              name.c_str(),
              property.m_value_string);
          break;
        default:
          std::cerr << "Invalid socket type "
                    << static_cast<int>(property.m_type) << std::endl;
      }
    }

    delete node_instance_accessor;
  }
}

bool AnimGraphResource::SaveStateMachineResourceToFile(
    const char* filename) const {
  assert(false && "Not yet implemented");

  return false;
}

bool AnimGraphResource::LoadStateMachineResourceFromJson(
    nlohmann::json const& json_data) {
  assert(false && "Not yet implemented");

  return false;
}