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

//
// 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_type_name;

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

  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) {
  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 =
      AnimNodeDescriptorFactory(result.m_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_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;
}

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 LoadBlendTreeResourceFromJson(json_data);
  } 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::LoadBlendTreeResourceFromJson(nlohmann::json const& json_data) {
  m_blend_tree_resource.Reset();
  m_type = "BlendTree";
  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);
    m_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 = m_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 = m_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);
    m_blend_tree_resource.m_connections.push_back(connection);
  }

  return true;
}

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

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

  return false;
}

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

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

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

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

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

  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_type != "BlendTree") {
    std::cerr << "Invalid AnimGraphResource. Expected type 'BlendTree' but got '"
    << m_type << "'." << std::endl;
    return;
  }

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

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

void AnimGraphResource::CreateBlendTreeRuntimeNodeInstances(
    AnimGraphBlendTree& result) const {
  for (int i = 0; i < m_blend_tree_resource.m_nodes.size(); i++) {
    const AnimNodeResource& node_resource = m_blend_tree_resource.m_nodes[i];
    AnimNode* node = AnimNodeFactory(node_resource.m_type_name);
    node->m_name = node_resource.m_name;
    node->m_node_type_name = node_resource.m_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);
  }

  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_type_name,
        instance.m_nodes[i]);
  }

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