//
// Created by martin on 04.02.22.
//

#include "AnimGraphResource.h"

#include <fstream>

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

namespace AniGraph {

using json = nlohmann::json;

//
// 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_value.ptr != nullptr) {
    if (socket.m_type == SocketType::SocketTypeBool) {
      result["value"] = *reinterpret_cast<bool*>(socket.m_value.ptr);
    } else if (socket.m_type == SocketType::SocketTypeAnimation) {
    } else if (socket.m_type == SocketType::SocketTypeFloat) {
      result["value"] = *reinterpret_cast<float*>(socket.m_value.ptr);
    } else if (socket.m_type == SocketType::SocketTypeVec3) {
      Vec3& vec3 = *reinterpret_cast<Vec3*>(socket.m_value.ptr);
      result["value"][0] = vec3[0];
      result["value"][1] = vec3[1];
      result["value"][2] = vec3[2];
    } else if (socket.m_type == SocketType::SocketTypeQuat) {
      Quat& quat = *reinterpret_cast<Quat*>(socket.m_value.ptr);
      result["value"][0] = quat[0];
      result["value"][1] = quat[1];
      result["value"][2] = quat[2];
      result["value"][3] = quat[3];
    } else if (socket.m_type == SocketType::SocketTypeString) {
      result["value"] = *reinterpret_cast<std::string*>(socket.m_value.ptr);
    } 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_value.ptr = nullptr;
  result.m_name = json_data["name"];

  std::string type_string = json_data["type"];

  if (type_string == "Bool") {
    result.m_type = SocketType::SocketTypeBool;
    result.m_type_size = sizeof(bool);
  } else if (type_string == "Animation") {
    result.m_type = SocketType::SocketTypeAnimation;
    result.m_type_size = sizeof(AnimData);
  } else if (type_string == "Float") {
    result.m_type = SocketType::SocketTypeFloat;
    result.m_type_size = sizeof(float);
  } else if (type_string == "Vec3") {
    result.m_type = SocketType::SocketTypeVec3;
    result.m_type_size = sizeof(Vec3);
  } else if (type_string == "Quat") {
    result.m_type = SocketType::SocketTypeQuat;
    result.m_type_size = sizeof(Quat);
  } else if (type_string == "String") {
    result.m_type = SocketType::SocketTypeString;
    result.m_type_size = sizeof(std::string);
  } else {
    std::cerr << "Invalid socket type '" << type_string << "'." << std::endl;
  }

  return result;
}

//
// AnimGraphNode <-> json
//
json sAnimGraphNodeToJson(const AnimNodeResource& node) {
  json result;

  result["name"] = node.m_name;
  result["type"] = "AnimNodeResource";
  result["node_type"] = node.m_type_name;

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

  for (size_t j = 0, n = node.m_socket_accessor->m_properties.size(); j < n;
       j++) {
    Socket& property = node.m_socket_accessor->m_properties[j];
    result["properties"][property.m_name] = sSocketToJson(property);
  }

  return result;
}

AnimNodeResource sAnimGraphNodeFromJson(const json& json_node) {
  AnimNodeResource result;

  result.m_name = json_node["name"];
  result.m_type_name = json_node["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_type_name);
  result.m_socket_accessor =
      AnimNodeAccessorFactory(result.m_type_name, result.m_anim_node);

  for (size_t j = 0, n = result.m_socket_accessor->m_properties.size(); j < n;
       j++) {
    Socket& property = result.m_socket_accessor->m_properties[j];
    json json_property = json_node["properties"][property.m_name];

    if (sSocketTypeToStr(property.m_type) == json_property["type"]) {
      if (property.m_type == SocketType::SocketTypeBool) {
        result.m_socket_accessor->SetProperty<bool>(
            property.m_name,
            json_property["value"]);
      } else if (property.m_type == SocketType::SocketTypeAnimation) {
      } else if (property.m_type == SocketType::SocketTypeFloat) {
        result.m_socket_accessor->SetProperty<float>(
            property.m_name,
            json_property["value"]);
      } else if (property.m_type == SocketType::SocketTypeVec3) {
        Vec3* property_vec3 = reinterpret_cast<Vec3*>(property.m_value.ptr);
        (*property_vec3)[0] = json_property["value"][0];
        (*property_vec3)[1] = json_property["value"][1];
        (*property_vec3)[2] = json_property["value"][2];
      } else if (property.m_type == SocketType::SocketTypeQuat) {
        Quat* property_quat = reinterpret_cast<Quat*>(property.m_value.ptr);
        (*property_quat)[0] = json_property["value"][0];
        (*property_quat)[1] = json_property["value"][1];
        (*property_quat)[2] = json_property["value"][2];
        (*property_quat)[3] = json_property["value"][3];
      } else if (property.m_type == SocketType::SocketTypeString) {
        result.m_socket_accessor->SetProperty<std::string>(
            property.m_name,
            json_property["value"]);
      } else {
        std::cerr << "Invalid type for property '" << property.m_name
                  << "'. Cannot parse json to type '"
                  << static_cast<int>(property.m_type) << std::endl;
        break;
      }
    } else {
      std::cerr << "Invalid type for property '" << property.m_name
                << "': expected " << sSocketTypeToStr(property.m_type)
                << " but got " << json_property["type"] << std::endl;
    }
  }

  return result;
}

//
// AnimGraphConnection <-> Json
//
json sAnimGraphConnectionToJson(const AnimGraphResource& graph_resource, const AnimGraphConnection& connection) {
  json result;

  result["type"] = "AnimGraphConnection";

  const AnimNodeResource* source_node = connection.m_source_node;
  result["source_node_index"] = graph_resource.getNodeIndex(*source_node);
  result["source_socket_index"] = source_node->m_socket_accessor->GetOutputIndex(connection.m_source_socket->m_name);

  const AnimNodeResource* target_node = connection.m_source_node;
  result["source_node_index"] = graph_resource.getNodeIndex(*target_node);
  result["source_socket_index"] = target_node->m_socket_accessor->GetInputIndex(connection.m_target_node->m_name);

  return result;
}

AnimGraphConnection sAnimGraphConnectionFromJson(const AnimGraphResource& graph_resource, const json& json_node) {
  AnimGraphConnection connection;

  int source_node_index = json_node["source_node_index"];
  connection.m_source_node = &graph_resource.m_nodes[source_node_index];
  int source_socket_index = json_node["source_socket_index"];
  connection.m_source_socket = &connection.m_source_node->m_socket_accessor->m_outputs[source_socket_index];

  
  int target_node_index = json_node["target_node_index"];
  connection.m_target_node = &graph_resource.m_nodes[target_node_index];
  
  int target_socket_index = json_node["target_socket_index"];
  connection.m_target_socket = &connection.m_target_node->m_socket_accessor->m_outputs[target_socket_index];

  return connection;
}

void AnimGraphResource::clear() {
  m_name = "";

  clearNodes();
  m_connections.clear();

  initGraphConnectors();
}

void AnimGraphResource::clearNodes() {
  for (size_t i = 0; i < m_nodes.size(); i++) {
    delete m_nodes[i].m_socket_accessor;
    m_nodes[i].m_socket_accessor = nullptr;
    delete m_nodes[i].m_anim_node;
    m_nodes[i].m_anim_node = nullptr;
  }
  m_nodes.clear();
}

void AnimGraphResource::initGraphConnectors() {
  m_nodes.push_back(AnimNodeResourceFactory("BlendTree"));
  m_nodes[0].m_name = "Outputs";
  m_nodes.push_back(AnimNodeResourceFactory("BlendTree"));
  m_nodes[1].m_name = "Inputs";
}

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

  result["name"] = m_name;
  result["type"] = "AnimGraphResource";

  for (size_t i = 0; i < m_nodes.size(); i++) {
    const AnimNodeResource& node = m_nodes[i];
    result["nodes"][i] = sAnimGraphNodeToJson(node);
  }

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

  // Graph inputs and outputs
  {
    const AnimNodeResource& graph_output_node = 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 = 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]);
    }
  }

  std::ofstream output_file;
  output_file.open(filename);
  output_file << to_string(result) << std::endl;
  output_file.close();

  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;
  }

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

  clear();
  clearNodes();

  m_name = json_data["name"];
  
  // Load nodes
  for (size_t i = 0; i < json_data["nodes"].size(); 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);
    m_nodes.push_back(node);
  }
  
  // Setup graph inputs and outputs
  const json& graph_outputs = json_data["nodes"][0]["inputs"];
  for (size_t i = 0; i < graph_outputs.size(); i++) {
    AnimNodeResource& graph_node = m_nodes[0];
    graph_node.m_socket_accessor->m_inputs.push_back(
        sJsonToSocket(graph_outputs[i]));
  }
  
  const json& graph_inputs = json_data["nodes"][1]["outputs"];
  for (size_t i = 0; i < graph_inputs.size(); i++) {
    AnimNodeResource& graph_node = m_nodes[1];
    graph_node.m_socket_accessor->m_outputs.push_back(
        sJsonToSocket(graph_inputs[i]));
  }

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

    AnimGraphConnection connection =
        sAnimGraphConnectionFromJson(*this, json_connection);
    m_connections.push_back(connection);
  }

  return true;
}

void AnimGraph::updateOrderedNodes() {
  std::vector<int> node_index_stack;
  node_index_stack.push_back(0);

  m_eval_ordered_nodes.clear();

  while (node_index_stack.size() > 0) {
    std::vector<NodeInput>& node_inputs =
        m_node_inputs[node_index_stack.back()];
    node_index_stack.pop_back();

    for (size_t i = 0, n = node_inputs.size(); i < n; i++) {
      AnimNode* input_node = node_inputs[i].m_node;
      if (input_node == nullptr) {
        continue;
      }

      int input_node_index = input_node->m_index;
      bool is_node_processed = false;
      for (size_t j = 0, m = m_eval_ordered_nodes.size(); j < m; j++) {
        if (m_eval_ordered_nodes[j] == input_node) {
          is_node_processed = true;
          break;
        }
      }

      if (is_node_processed) {
        continue;
      }

      m_eval_ordered_nodes.push_back(input_node);
      node_index_stack.push_back(input_node_index);
    }
  }
}

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<NodeInput> graph_output_inputs = m_node_inputs[0];
  for (size_t i = 0, n = graph_output_inputs.size(); i < n; i++) {
    const NodeInput& graph_input = graph_output_inputs[i];
    AnimNode* node = graph_input.m_node;
    if (node != nullptr) {
      node->m_state = AnimNodeEvalState::Activated;
    }
  }

  for (size_t i = 0, n = m_eval_ordered_nodes.size(); i < n; i++) {
    AnimNode* node = m_eval_ordered_nodes[i];
    if (checkIsNodeActive(node)) {
      int node_index = node->m_index;
      node->MarkActiveInputs(m_node_inputs[node_index]);

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

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_inputs[node_index]);
  }
}

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

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

    int node_index = node->m_index;
    const std::vector<NodeInput> node_inputs =
        m_node_inputs[node_index];
    float node_time_now = node->m_time_now;
    float node_time_last = node->m_time_last;

    for (size_t i = 0, n = node_inputs.size(); i < n; i++) {
      AnimNode* input_node = node_inputs[i].m_node;

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

void AnimGraph::evaluate() {
  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::Deactivated) {
      continue;
    }

    node->Evaluate();
  }
}

void* AnimGraph::getOutput(const std::string& name) const {
  Socket* socket = m_socket_accessor->FindInputSocket(name);
  if (socket == nullptr) {
    return nullptr;
  }

  return socket->m_value.ptr;
}

void* AnimGraph::getInput(const std::string& name) const {
  Socket* socket = m_socket_accessor->FindOutputSocket(name);
  if (socket == nullptr) {
    return nullptr;
  }

  return *(socket->m_value.ptr_ptr);
}

}  // namespace AniGraph