//
// Created by martin on 11.02.22.
//

#include "AnimGraphEditor.h"

#include <sstream>

#include "3rdparty/imgui-node-editor/imgui_node_editor.h"
#include "AnimGraphEditor/utilities/builders.h"
#include "AnimGraphEditor/utilities/drawing.h"
#include "AnimGraphEditor/utilities/widgets.h"
#include "SkinnedMesh.h"
#include "imgui.h"
#include "imnodes.h"
#include "misc/cpp/imgui_stdlib.h"
#include "src/AnimGraph/AnimGraphResource.h"

struct EditorState {
  AnimGraphResource* rootGraphResource = nullptr;

  std::vector<AnimGraphResource*> hierarchyStack;
  size_t hierarchyStackIndex = 0;

  bool isGraphLoadedThisFrame = false;
  ImVec2 mousePopupStart = {};
  char statusLine[1024] = {};
};

struct NodeConnectionDebugState {
  struct SocketInfo {
    ax::NodeEditor::PinId pin = {};
    int nodeId = {};
    const AnimNodeResource* nodeResource = {};
    int socketId = {};
    const Socket* socket = {};
  };

  SocketInfo sourceSocket = {};
  SocketInfo targetSocket = {};

  void Reset() {
    sourceSocket = {};
    targetSocket = {};
  }
};

static EditorState sEditorState;
static NodeConnectionDebugState sNodeConnectionDebugState;

constexpr int cPinIconSize = 24;

ImNodesPinShape sGetSocketShapeFromSocketType(const SocketType& socket_type) {
  switch (socket_type) {
    case SocketType::SocketTypeAnimation:
      return ImNodesPinShape_QuadFilled;
    case SocketType::SocketTypeInt:
      return ImNodesPinShape_CircleFilled;
    case SocketType::SocketTypeFloat:
      return ImNodesPinShape_CircleFilled;
    case SocketType::SocketTypeVec3:
      return ImNodesPinShape_TriangleFilled;
    case SocketType::SocketTypeQuat:
      return ImNodesPinShape_Triangle;
    case SocketType::SocketTypeBool:
      return ImNodesPinShape_Circle;
    default:
      break;
  }

  return ImNodesPinShape_Quad;
}

void DrawSocketIcon(const SocketType& socket_type, bool is_connected = false) {
  ax::Drawing::IconType iconType;
  ImColor color = ImColor(255, 255, 255);
  color.Value.w = 1;
  switch (socket_type) {
    case SocketType::SocketTypeAnimation:
      iconType = ax::Drawing::IconType::Flow;
      break;
    case SocketType::SocketTypeInt:
      iconType = ax::Drawing::IconType::Circle;
      break;
    case SocketType::SocketTypeFloat:
      iconType = ax::Drawing::IconType::Circle;
      break;
    case SocketType::SocketTypeVec3:
      iconType = ax::Drawing::IconType::Circle;
      break;
    case SocketType::SocketTypeQuat:
      iconType = ax::Drawing::IconType::Circle;
      break;
    case SocketType::SocketTypeBool:
      iconType = ax::Drawing::IconType::Circle;
      break;
    default:
      break;
  }

  ax::Widgets::Icon(
      ImVec2(
          static_cast<float>(cPinIconSize),
          static_cast<float>(cPinIconSize)),
      iconType,
      is_connected,
      color,
      ImColor(32, 32, 32, 1));
}

bool NodeSocketEditor(Socket& socket) {
  bool modified = false;
  int mode_current = static_cast<int>(socket.m_type);
  ImGui::Columns(3);
  if (ImGui::Combo(
          "Type",
          &mode_current,
          SocketTypeNames,
          sizeof(SocketTypeNames) / sizeof(char*))) {
    socket.m_type = static_cast<SocketType>(mode_current);

    modified = true;
  }

  ImGui::Columns(2);
  if (ImGui::InputText("Name", &socket.m_name)) {
    modified = true;
  }

  ImGui::Columns(1);

  return modified;
}

void SyncTrackEditor(SyncTrack* sync_track) {
  ImGui::SliderFloat("duration", &sync_track->m_duration, 0.001f, 10.f);

  ImGui::Text("Marker");
  ImGui::SameLine();
  ImGui::Text("%d", sync_track->m_num_intervals);
  ImGui::SameLine();
  if (ImGui::Button("+")) {
    if (sync_track->m_num_intervals < cSyncTrackMaxIntervals) {
      sync_track->m_num_intervals++;
    }
  }
  ImGui::SameLine();
  if (ImGui::Button("-")) {
    if (sync_track->m_num_intervals > 0) {
      sync_track->m_num_intervals--;
    }
  }

  ImGui::Text("Marker:");
  for (int i = 0; i < sync_track->m_num_intervals; i++) {
    ImGui::Text("%2d:", i);
    ImGui::SameLine();
    std::ostringstream marker_stream;
    marker_stream << i;
    ImGui::SliderFloat(
        marker_stream.str().c_str(),
        &sync_track->m_sync_markers[i],
        0.f,
        1.f);
  }

  if (ImGui::Button("Update Intervals")) {
    sync_track->CalcIntervals();
  }
}

void SkinnedMeshWidget(SkinnedMesh* skinned_mesh) {
  if (ImGui::TreeNode("Bones")) {
    for (int i = 0; i < skinned_mesh->m_skeleton.num_joints(); i++) {
      ImGui::Text("%s", skinned_mesh->m_skeleton.joint_names()[i]);
    }
    ImGui::TreePop();
  }

  ImGui::Text("Animations");

  const char* items[255] = {0};
  static int selected = -1;
  for (int i = 0; i < skinned_mesh->m_animations.size(); i++) {
    items[i] = skinned_mesh->m_animation_names[i].c_str();
  }

  ImGui::Combo(
      "Animation",
      &selected,
      items,
      skinned_mesh->m_animations.size());

  ImGui::Text("Sync Track");
  if (selected >= 0 && selected < skinned_mesh->m_animations.size()) {
    SyncTrackEditor(&skinned_mesh->m_animation_sync_track[selected]);
    skinned_mesh->m_override_anim = selected;

    ImGui::Checkbox("Override Animation", &skinned_mesh->m_sync_track_override);
    if (skinned_mesh->m_sync_track_override) {
      ImGui::SliderFloat("Ratio", &skinned_mesh->m_override_ratio, 0.f, 1.f);

      ozz::animation::SamplingJob sampling_job;
      sampling_job.animation = skinned_mesh->m_animations[selected];
      sampling_job.context = &skinned_mesh->m_sampling_context;
      sampling_job.ratio = skinned_mesh->m_override_ratio;
      sampling_job.output = make_span(skinned_mesh->m_local_matrices);
      if (!sampling_job.Run()) {
        ozz::log::Err() << "Error sampling animation." << std::endl;
      }
    }
  }
}

void AnimGraphEditorRenderSidebar(
    BlendTreeResource& blend_tree_resource,
    AnimNodeResource* node_resource) {
  ImGui::Text(
      "[%s (%2.2f, %2.2f)]",
      node_resource->m_node_type_name.c_str(),
      node_resource->m_position[0],
      node_resource->m_position[1]);

  char node_name_buffer[256];
  memset(node_name_buffer, 0, sizeof(node_name_buffer));
  strncpy(
      node_name_buffer,
      node_resource->m_name.c_str(),
      std::min(node_resource->m_name.size(), sizeof(node_name_buffer)));

  if (ImGui::InputText("Name", node_name_buffer, sizeof(node_name_buffer))) {
    node_resource->m_name = node_name_buffer;
  }

  int num_properties = 0;
  if (node_resource->m_virtual_socket_accessor != nullptr) {
    num_properties =
        node_resource->m_virtual_socket_accessor->m_properties.size();
  }
  for (int i = 0; i < num_properties; i++) {
    Socket& property =
        node_resource->m_virtual_socket_accessor->m_properties[i];
    if (property.m_type == SocketType::SocketTypeInt) {
      ImGui::InputInt(
          property.m_name.c_str(),
          reinterpret_cast<int*>(&property.m_value.int_value),
          1);
    } else if (property.m_type == SocketType::SocketTypeFloat) {
      ImGui::SliderFloat(
          property.m_name.c_str(),
          reinterpret_cast<float*>(&property.m_value.float_value),
          -100.f,
          100.f);
    } else if (property.m_type == SocketType::SocketTypeBool) {
      bool flag_value = property.GetValue<bool>();
      if (ImGui::Checkbox(property.m_name.c_str(), &flag_value)) {
        property.SetValue(flag_value);
      }
    } else if (property.m_type == SocketType::SocketTypeString) {
      char string_buf[1024];
      memset(string_buf, '\0', sizeof(string_buf));
      memcpy(
          string_buf,
          property.m_value_string.c_str(),
          std::min(
              static_cast<size_t>(1024),
              property.m_value_string.size() + 1));
      if (ImGui::InputText(
              property.m_name.c_str(),
              string_buf,
              sizeof(string_buf))) {
        property.m_value_string = string_buf;
      }
    }
  }

  if (node_resource == blend_tree_resource.GetGraphOutputNode()) {
    ImGui::Text("Outputs");

    // Graph outputs are the inputs of the output node!
    std::vector<Socket>& outputs =
        node_resource->m_virtual_socket_accessor->m_inputs;

    std::vector<Socket>::iterator iter = outputs.begin();
    ImGui::PushStyleVar(ImGuiStyleVar_FramePadding, ImVec2(2, 2));
    if (ImGui::BeginTable(
            "split",
            2,
            ImGuiTableFlags_BordersOuter | ImGuiTableFlags_Resizable)) {
      while (iter != outputs.end()) {
        Socket& socket = *iter;
        ImGui::PushID(&socket);

        ImGui::TableNextRow();
        ImGui::TableSetColumnIndex(0);
        ImGui::AlignTextToFramePadding();
        bool node_open = ImGui::TreeNode(socket.m_name.c_str());
        ImGui::TableSetColumnIndex(1);

        if (node_open) {
          int mode_current = static_cast<int>(socket.m_type);
          if (ImGui::Combo(
                  "Type",
                  &mode_current,
                  SocketTypeNames,
                  sizeof(SocketTypeNames) / sizeof(char*))) {
            socket.m_type = static_cast<SocketType>(mode_current);
          }
          ImGui::TreePop();
        }

        ImGui::PopID();
        iter++;
      }
      ImGui::EndTable();
    }
    ImGui::PopStyleVar();

    if (ImGui::Button("+")) {
      AnimGraphResource* current_graph_resource =
          sEditorState.hierarchyStack[sEditorState.hierarchyStackIndex];
      current_graph_resource->RegisterBlendTreeOutputSocket<float>(
          "GraphFloatOutput"
          + std::to_string(current_graph_resource->m_blend_tree_resource
                               .GetGraphOutputNode()
                               ->m_virtual_socket_accessor->m_inputs.size()));
    }
  }

  if (node_resource == blend_tree_resource.GetGraphInputNode()) {
    ImGui::Text("Inputs");

    // Graph inputs are the outputs of the input node!
    std::vector<Socket>& inputs =
        node_resource->m_virtual_socket_accessor->m_outputs;

    std::vector<Socket>::iterator iter = inputs.begin();
    while (iter != inputs.end()) {
      Socket& input = *iter;
      ImGui::PushID(&input);
      if (NodeSocketEditor(input)) {
        AnimGraphResource* current_graph_resource =
            sEditorState.hierarchyStack[sEditorState.hierarchyStackIndex];
        current_graph_resource->m_virtual_socket_accessor->m_inputs = inputs;
      }
      if (ImGui::Button("X")) {
        blend_tree_resource.RemoveConnectionsForSocket(node_resource, input);
        iter = inputs.erase(iter);
      } else {
        iter++;
      }
      ImGui::PopID();
    }

    if (ImGui::Button("+")) {
      AnimGraphResource* current_graph_resource =
          sEditorState.hierarchyStack[sEditorState.hierarchyStackIndex];
      current_graph_resource->RegisterBlendTreeInputSocket<float>(
          "GraphFloatInput"
          + std::to_string(
              current_graph_resource->m_blend_tree_resource.GetGraphInputNode()
                  ->m_virtual_socket_accessor->m_outputs.size()));
    }
  }
}

void AnimGraphEditorResetHierarchyStack() {
  sEditorState.hierarchyStack.clear();
  sEditorState.hierarchyStack.push_back(sEditorState.rootGraphResource);
  sEditorState.hierarchyStack[sEditorState.hierarchyStackIndex] =
      sEditorState.hierarchyStack.back();
  sEditorState.hierarchyStackIndex = 0;
}

void AnimGraphEditorClear() {
  if (ax::NodeEditor::GetCurrentEditor() != nullptr) {
    ax::NodeEditor::ClearSelection();
  }

  if (sEditorState.rootGraphResource) {
    delete sEditorState.rootGraphResource;
  }

  sEditorState.rootGraphResource =
      dynamic_cast<AnimGraphResource*>(AnimNodeResourceFactory("BlendTree"));
  sEditorState.rootGraphResource->m_name = "Root";
  sEditorState.isGraphLoadedThisFrame = true;

  AnimGraphEditorResetHierarchyStack();
}

void BlendTreeEditorNodePopup() {
  const bool open_popup = ImGui::IsMouseReleased(ImGuiMouseButton_Right);

  if (open_popup && ImGui::IsWindowHovered()) {
    ax::NodeEditor::Suspend();
    ImGui::OpenPopup("add node");
    ax::NodeEditor::Resume();
    sEditorState.mousePopupStart = ImGui::GetMousePos();
  }

  ax::NodeEditor::Suspend();
  ImGui::PushStyleVar(ImGuiStyleVar_WindowPadding, ImVec2(8.f, 8.f));
  if (ImGui::BeginPopup("add node")) {
    std::string node_type_name = "";
    if (ImGui::MenuItem("AnimSampler")) {
      node_type_name = "AnimSampler";
    }

    if (ImGui::MenuItem("Blend2")) {
      node_type_name = "Blend2";
    }

    if (ImGui::MenuItem("SpeedScale")) {
      node_type_name = "SpeedScale";
    }

    if (ImGui::MenuItem("LockTranslationNode")) {
      node_type_name = "LockTranslationNode";
    }

    if (ImGui::MenuItem("MathAddNode")) {
      node_type_name = "MathAddNode";
    }

    if (ImGui::MenuItem("MathFloatToVec3Node")) {
      node_type_name = "MathFloatToVec3Node";
    }

    if (ImGui::MenuItem("ConstScalarNode")) {
      node_type_name = "ConstScalarNode";
    }

    if (ImGui::MenuItem("BlendTree")) {
      node_type_name = "BlendTree";
    }

    if (!node_type_name.empty()) {
      AnimNodeResource* node_resource = AnimNodeResourceFactory(node_type_name);
      ax::NodeEditor::SetNodePosition(
          ax::NodeEditor::NodeId(node_resource),
          sEditorState.mousePopupStart);
      sEditorState.hierarchyStack[sEditorState.hierarchyStackIndex]
          ->m_blend_tree_resource.AddNode(node_resource);
    }

    ImGui::EndPopup();
  }
  ImGui::PopStyleVar();
  ax::NodeEditor::Resume();
}

void AnimGraphEditorMenuBar() {
  ImGui::BeginMenuBar();
  if (ImGui::Button("Save")) {
    sEditorState.rootGraphResource->SaveToFile("editor_graph.json");
  }
  if (ImGui::Button("Load")) {
    AnimGraphEditorClear();

    delete sEditorState.rootGraphResource;

    sEditorState.rootGraphResource =
        AnimGraphResource::CreateFromFile("editor_graph.json");

    AnimGraphEditorResetHierarchyStack();
  }
  if (ImGui::Button("Clear")) {
    AnimGraphEditorClear();
  }
  if (ImGui::Button("Content")) {
    ax::NodeEditor::NavigateToContent();
  }
  char graph_name_buffer[256];
  memset(graph_name_buffer, 0, sizeof(graph_name_buffer));
  strncpy(
      graph_name_buffer,
      sEditorState.hierarchyStack[sEditorState.hierarchyStackIndex]
          ->m_name.c_str(),
      sizeof(graph_name_buffer));
  if (ImGui::InputText("Name", graph_name_buffer, sizeof(graph_name_buffer))) {
    sEditorState.hierarchyStack[sEditorState.hierarchyStackIndex]->m_name =
        graph_name_buffer;
  }

  ImGui::EndMenuBar();
}

void AnimGraphEditorBreadcrumbNavigation() {
  for (size_t i = 0, n = sEditorState.hierarchyStack.size(); i < n; i++) {
    AnimGraphResource* graph_resource =
        dynamic_cast<AnimGraphResource*>(sEditorState.hierarchyStack[i]);
    ImGui::PushID(graph_resource);

    bool highlight_button = i == sEditorState.hierarchyStackIndex;

    if (highlight_button) {
      ImGui::PushStyleColor(
          ImGuiCol_Button,
          (ImVec4)ImColor::HSV(1. / 7.0f, 0.6f, 0.6f));
    }

    if (ImGui::Button(graph_resource->m_name.c_str())) {
      sEditorState.hierarchyStackIndex = i;
    }

    if (highlight_button) {
      ImGui::PopStyleColor(1);
    }

    ImGui::PopID();

    if (i < n - 1) {
      ImGui::SameLine();
    }
  }
}

void HandleConnectionCreation(BlendTreeResource& current_blend_tree) {
  if (ax::NodeEditor::BeginCreate()) {
    ax::NodeEditor::PinId input_pin_id, output_pin_id;
    if (ax::NodeEditor::QueryNewLink(&input_pin_id, &output_pin_id)) {
      if (input_pin_id == output_pin_id) {
        ax::NodeEditor::RejectNewItem();
        ax::NodeEditor::EndCreate();
        return;
      }

      ax::NodeEditor::PinId source_pin = input_pin_id;
      ax::NodeEditor::PinId target_pin = output_pin_id;

      int source_node_index = -1;
      int source_node_socket_index = -1;
      const AnimNodeResource* source_node = nullptr;
      const Socket* source_socket = nullptr;

      if (input_pin_id) {
        if (IsPinInput(input_pin_id.Get())) {
          source_pin = input_pin_id;
          target_pin = output_pin_id;
        } else {
          target_pin = input_pin_id;
          source_pin = output_pin_id;
        }
      }

      if (output_pin_id) {
        if (IsPinOutput(output_pin_id.Get())) {
          target_pin = output_pin_id;
          source_pin = input_pin_id;
        } else {
          source_pin = output_pin_id;
          target_pin = input_pin_id;
        }
      }

      if (!source_pin.Invalid) {
        OutputPinIdToNodeIndexAndSocketIndex(
            source_pin.Get(),
            &source_node_index,
            &source_node_socket_index);

        source_node = current_blend_tree.GetNode(source_node_index);

        if (source_node != nullptr) {
          if (source_node->m_virtual_socket_accessor->m_outputs.size()
              < source_node_socket_index) {
            source_node_socket_index = -1;
          } else {
            source_socket = current_blend_tree.GetNodeOutputSocketByIndex(
                source_node,
                source_node_socket_index);
          }
        }
      }

      int target_node_index = -1;
      int target_node_socket_index = -1;
      const AnimNodeResource* target_node = nullptr;
      const Socket* target_socket = nullptr;

      if (!target_pin.Invalid) {
        InputPinIdToNodeIndexAndSocketIndex(
            target_pin.Get(),
            &target_node_index,
            &target_node_socket_index);

        target_node = current_blend_tree.GetNode(target_node_index);

        if (target_node != nullptr) {
          if (target_node->m_virtual_socket_accessor->m_inputs.size()
              < target_node_socket_index) {
            target_node_socket_index = -1;
          } else {
            target_socket = current_blend_tree.GetNodeInputSocketByIndex(
                target_node,
                target_node_socket_index);
          }
        }
      }

      sNodeConnectionDebugState.sourceSocket.pin = source_pin;
      sNodeConnectionDebugState.sourceSocket.nodeId = source_node_index;
      sNodeConnectionDebugState.sourceSocket.nodeResource = source_node;
      sNodeConnectionDebugState.sourceSocket.socketId =
          source_node_socket_index;
      sNodeConnectionDebugState.sourceSocket.socket = source_socket;

      sNodeConnectionDebugState.targetSocket.pin = target_pin;
      sNodeConnectionDebugState.targetSocket.nodeId = target_node_index;
      sNodeConnectionDebugState.targetSocket.nodeResource = target_node;
      sNodeConnectionDebugState.targetSocket.socketId =
          target_node_socket_index;
      sNodeConnectionDebugState.targetSocket.socket = target_socket;

      if (!source_pin.Invalid && !target_pin.Invalid) {
        if (source_socket == nullptr || target_socket == nullptr
            || !current_blend_tree.IsConnectionValid(
                source_node,
                source_socket->m_name,
                target_node,
                target_socket->m_name)) {
          ax::NodeEditor::RejectNewItem();
        } else if (ax::NodeEditor::AcceptNewItem()) {
          current_blend_tree.ConnectSockets(
              source_node,
              source_socket->m_name,
              target_node,
              target_socket->m_name);
        }
      }
    }
  }
  ax::NodeEditor::EndCreate();
}

void BlendTreeRenderNodes(
    BlendTreeResource& current_blend_tree,
    ax::NodeEditor::Utilities::BlueprintNodeBuilder& builder) {
  for (size_t node_index = 0, n = current_blend_tree.GetNumNodes();
       node_index < n;
       node_index++) {
    AnimNodeResource* node_resource = current_blend_tree.GetNode(node_index);

    ax::NodeEditor::NodeId node_id(node_resource);

    builder.Begin(node_id);

    if (sEditorState.isGraphLoadedThisFrame) {
      ax::NodeEditor::SetNodePosition(
          node_id,
          ImVec2(node_resource->m_position[0], node_resource->m_position[1]));
    }

    builder.Header();
    if (node_resource->m_name != "") {
      ImGui::Text("%s", node_resource->m_name.c_str());
    } else {
      ImGui::Text("%s", node_resource->m_node_type_name.c_str());
    }
    ImGui::Spring(0);
    builder.EndHeader();

    // Inputs
    std::vector<Socket> node_inputs =
        current_blend_tree.GetNodeInputSockets(node_resource);
    for (size_t j = 0, ni = node_inputs.size(); j < ni; j++) {
      Socket& socket = node_inputs[j];
      ax::NodeEditor::PinId input_pin = NodeIndexAndSocketIndexToInputPinId(
          static_cast<int>(node_index),
          static_cast<int>(j));
      builder.Input(input_pin);

      assert(!input_pin.Invalid);

      DrawSocketIcon(
          socket.m_type,
          current_blend_tree.IsSocketConnected(node_resource, socket.m_name));
      ImGui::Spring(0);

      //ImGui::PushItemWidth(100.0f);
      ImGui::Text("%s", socket.m_name.c_str());
      //ImGui::PopItemWidth();
      builder.EndInput();
    }

    // Outputs
    std::vector<Socket> node_outputs =
        current_blend_tree.GetNodeOutputSockets(node_resource);
    for (size_t j = 0, ni = node_outputs.size(); j < ni; j++) {
      Socket& socket = node_outputs[j];
      builder.Output(NodeIndexAndSocketIndexToOutputPinId(
          static_cast<int>(node_index),
          static_cast<int>(j)));

      ImGui::Text("%s", socket.m_name.c_str());

      ImGui::Spring(0);
      DrawSocketIcon(
          socket.m_type,
          current_blend_tree.IsSocketConnected(node_resource, socket.m_name));

      builder.EndOutput();
    }

    ImVec2 node_position = ax::NodeEditor::GetNodePosition(node_id);
    node_resource->m_position[0] = node_position.x;
    node_resource->m_position[1] = node_position.y;

    builder.End();
  }
}

void BlendTreeRenderConnections(BlendTreeResource& current_blend_tree) {
  for (size_t connection_id = 0, n = current_blend_tree.GetNumConnections();
       connection_id < n;
       connection_id++) {
    const BlendTreeConnectionResource* connection_resource =
        current_blend_tree.GetConnection(connection_id);

    const AnimNodeResource* source_node_resource =
        current_blend_tree.GetNode(connection_resource->source_node_index);
    int source_socket_index =
        source_node_resource->m_virtual_socket_accessor->GetOutputIndex(
            connection_resource->source_socket_name.c_str());

    const AnimNodeResource* target_node_resource =
        current_blend_tree.GetNode(connection_resource->target_node_index);
    int target_socket_index =
        target_node_resource->m_virtual_socket_accessor->GetInputIndex(
            connection_resource->target_socket_name.c_str());

    int source_socket_pin_id = NodeIndexAndSocketIndexToOutputPinId(
        static_cast<int>(connection_resource->source_node_index),
        source_socket_index);
    int target_socket_pin_id = NodeIndexAndSocketIndexToInputPinId(
        static_cast<int>(connection_resource->target_node_index),
        target_socket_index);

    ax::NodeEditor::Link(
        ax::NodeEditor::LinkId(connection_resource),
        source_socket_pin_id,
        target_socket_pin_id);
  }
}
void AnimGraphEditorDebugWidget() {
  ImGui::Begin("Connection Debug Panel");
  ImGui::BeginTable("Connection", 3);
  ImGui::TableNextRow();
  ImGui::TableSetColumnIndex(0);
  ImGui::Text("Pin");
  ImGui::TableNextColumn();
  ImGui::Text("%x", sNodeConnectionDebugState.sourceSocket.pin.AsPointer());
  ImGui::TableNextColumn();
  ImGui::Text("%x", sNodeConnectionDebugState.targetSocket.pin.AsPointer());

  ImGui::TableNextRow();
  ImGui::TableSetColumnIndex(0);
  ImGui::Text("Node");
  ImGui::TableNextColumn();
  if (sNodeConnectionDebugState.sourceSocket.nodeResource) {
    ImGui::Text(
        "%s (%p)",
        sNodeConnectionDebugState.sourceSocket.nodeResource->m_node_type_name
            .c_str(),
        sNodeConnectionDebugState.sourceSocket.nodeResource);
  }
  ImGui::TableNextColumn();
  if (sNodeConnectionDebugState.targetSocket.nodeResource) {
    ImGui::Text(
        "%s (%p)",
        sNodeConnectionDebugState.targetSocket.nodeResource->m_node_type_name
            .c_str(),
        sNodeConnectionDebugState.targetSocket.nodeResource);
  }

  ImGui::TableNextRow();
  ImGui::TableSetColumnIndex(0);
  ImGui::Text("NodeId");
  ImGui::TableNextColumn();
  ImGui::Text("%d", sNodeConnectionDebugState.sourceSocket.nodeId);
  ImGui::TableNextColumn();
  ImGui::Text("%d", sNodeConnectionDebugState.targetSocket.nodeId);

  ImGui::TableNextRow();
  ImGui::TableSetColumnIndex(0);
  ImGui::Text("Socket");
  ImGui::TableNextColumn();
  if (sNodeConnectionDebugState.sourceSocket.socket) {
    ImGui::Text(
        "%s (%p)",
        sNodeConnectionDebugState.sourceSocket.socket->m_name.c_str(),
        sNodeConnectionDebugState.sourceSocket.socket);
  }

  ImGui::TableNextColumn();
  if (sNodeConnectionDebugState.targetSocket.socket) {
    ImGui::Text(
        "%s (%p)",
        sNodeConnectionDebugState.targetSocket.socket->m_name.c_str(),
        sNodeConnectionDebugState.targetSocket.socket);
  }

  ImGui::TableNextRow();
  ImGui::TableSetColumnIndex(0);
  ImGui::Text("SocketId");
  ImGui::TableNextColumn();
  ImGui::Text("%x", sNodeConnectionDebugState.sourceSocket.socketId);
  ImGui::TableNextColumn();
  ImGui::Text("%x", sNodeConnectionDebugState.targetSocket.socketId);

  ImGui::EndTable();

  ImGui::End();
}

void AnimGraphEditorUpdate(ax::NodeEditor::EditorContext* context) {
  sEditorState.statusLine[0] = '\0';
  sEditorState.statusLine[sizeof(sEditorState.statusLine) - 1] = '\0';

  sNodeConnectionDebugState.Reset();

  ax::NodeEditor::SetCurrentEditor(context);

  AnimGraphEditorMenuBar();

  AnimGraphEditorBreadcrumbNavigation();

  static ImGuiTableFlags flags =
      ImGuiTableFlags_SizingStretchSame | ImGuiTableFlags_Resizable;
  ImGui::BeginTable("GraphEditorWithSidebar", 2, flags);
  ImGui::TableNextRow();
  ImGui::TableSetColumnIndex(0);

  //
  // Node editor canvas
  //
  ImVec2 graph_size = ImGui::GetContentRegionAvail();
  graph_size.y -= 20;
  ax::NodeEditor::Begin("Graph Editor", graph_size);

  AnimGraphResource* current_graph =
      sEditorState.hierarchyStack[sEditorState.hierarchyStackIndex];
  BlendTreeResource& current_blend_tree = current_graph->m_blend_tree_resource;

  ax::NodeEditor::Utilities::BlueprintNodeBuilder builder;

  BlendTreeRenderNodes(current_blend_tree, builder);
  BlendTreeRenderConnections(current_blend_tree);

  HandleConnectionCreation(current_blend_tree);
  BlendTreeEditorNodePopup();

  ax::NodeEditor::End();

  ImGui::Text("Status: %s", sEditorState.statusLine);

  //
  // Sidebar
  //
  ImGui::TableSetColumnIndex(1);

  if (ax::NodeEditor::GetSelectedObjectCount() > 0) {
    ax::NodeEditor::NodeId selected_node_id = 0;

    ax::NodeEditor::GetSelectedNodes(&selected_node_id, 1);

    if (selected_node_id.Get() != 0) {
      AnimGraphEditorRenderSidebar(
          sEditorState.hierarchyStack[sEditorState.hierarchyStackIndex]
              ->m_blend_tree_resource,
          selected_node_id.AsPointer<AnimNodeResource>());
    }
  }

  ImGui::EndTable();

  AnimGraphEditorDebugWidget();

  // Clear flag, however it may be re-set further down when handling double
  // clicking into subgraphs.
  sEditorState.isGraphLoadedThisFrame = false;

  //
  // Handle double click into subgraphs
  //
  ax::NodeEditor::NodeId double_clicked_node_id =
      ax::NodeEditor::GetDoubleClickedNode();
  if (!double_clicked_node_id.Invalid) {
    AnimNodeResource* clicked_node_resource =
        double_clicked_node_id.AsPointer<AnimNodeResource>();

    if (clicked_node_resource != nullptr
        && clicked_node_resource->m_node_type_name == "BlendTree") {
      AnimGraphResource* clicked_graph_resource =
          dynamic_cast<AnimGraphResource*>(clicked_node_resource);

      assert(clicked_graph_resource != nullptr);

      if (sEditorState.hierarchyStack.size()
              > sEditorState.hierarchyStackIndex + 1
          && sEditorState.hierarchyStack[sEditorState.hierarchyStackIndex + 1]
                 == clicked_graph_resource) {
        sEditorState.hierarchyStackIndex++;
      } else {
        sEditorState.hierarchyStack.resize(
            sEditorState.hierarchyStackIndex + 1);
        sEditorState.hierarchyStack.push_back(clicked_graph_resource);
        sEditorState.hierarchyStackIndex++;
      }

      sEditorState.isGraphLoadedThisFrame = true;
      ax::NodeEditor::ClearSelection();
    }
  }

  ax::NodeEditor::LinkId hovered_link = ax::NodeEditor::GetHoveredLink();
  if (!hovered_link.Invalid) {
    BlendTreeConnectionResource* connection_resource =
        hovered_link.AsPointer<BlendTreeConnectionResource>();

    if (connection_resource && ImGui::IsKeyPressed(ImGuiKey_Delete)) {
      BlendTreeResource* blend_tree_resource =
          &sEditorState.hierarchyStack[sEditorState.hierarchyStackIndex]
               ->m_blend_tree_resource;

      blend_tree_resource->DisconnectSockets(
          blend_tree_resource->GetNode(connection_resource->source_node_index),
          connection_resource->source_socket_name,
          blend_tree_resource->GetNode(connection_resource->target_node_index),
          connection_resource->target_socket_name);

      ax::NodeEditor::DeleteLink(hovered_link);
    }
  }

  ax::NodeEditor::NodeId hovered_node = ax::NodeEditor::GetHoveredNode();
  if (!hovered_node.Invalid) {
    AnimNodeResource* node_resource =
        hovered_node.AsPointer<AnimNodeResource>();

    if (node_resource && ImGui::IsKeyPressed(ImGuiKey_Delete)) {
      AnimGraphResource* current_graph_resource =
          sEditorState.hierarchyStack[sEditorState.hierarchyStackIndex];

      current_graph_resource->m_blend_tree_resource.RemoveNodeConnections(
          node_resource);
      current_graph_resource->m_blend_tree_resource.RemoveNode(node_resource);
    }
  }

  ax::NodeEditor::SetCurrentEditor(nullptr);
}

void AnimGraphEditorGetRuntimeGraph(AnimGraphBlendTree& blend_tree) {
  sEditorState.rootGraphResource->CreateBlendTreeInstance(blend_tree);
}