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

#include "ozz/base/io/archive.h"
#include "ozz/base/io/stream.h"
#include "ozz/base/log.h"

#include "AnimGraph/AnimGraph.h"
#include "AnimGraph/AnimGraphEditor.h"
#include "AnimGraph/AnimGraphResource.h"
#include "catch.hpp"

bool load_skeleton (ozz::animation::Skeleton& skeleton, const char* filename) {
  assert(filename);
  ozz::io::File file(filename, "rb");
  if (!file.opened()) {
    ozz::log::Err() << "Failed to open skeleton file " << filename << "."
                    << std::endl;
    return false;
  }
  ozz::io::IArchive archive(&file);
  if (!archive.TestTag<ozz::animation::Skeleton>()) {
    ozz::log::Err() << "Failed to load skeleton instance from file " << filename
                    << "." << std::endl;
    return false;
  }

  // Once the tag is validated, reading cannot fail.
  archive >> skeleton;

  return true;
}

TEST_CASE("BasicGraph", "[AnimGraphResource]") {
  AnimGraphResource graph_resource;

  graph_resource.clear();
  graph_resource.m_name = "WalkRunBlendGraph";

  // Prepare graph inputs and outputs
  size_t walk_node_index =
      graph_resource.addNode(AnimNodeResourceFactory("AnimSampler"));
  size_t run_node_index =
      graph_resource.addNode(AnimNodeResourceFactory("AnimSampler"));
  size_t blend_node_index =
      graph_resource.addNode(AnimNodeResourceFactory("Blend2"));

  AnimNodeResource& walk_node = graph_resource.m_nodes[walk_node_index];
  walk_node.m_name = "WalkAnim";
  walk_node.m_socket_accessor->SetProperty("Filename", std::string("data/walk.anim.ozz"));
  AnimNodeResource& run_node = graph_resource.m_nodes[run_node_index];
  run_node.m_socket_accessor->SetProperty("Filename", std::string("data/run.anim.ozz"));
  run_node.m_name = "RunAnim";
  AnimNodeResource& blend_node = graph_resource.m_nodes[blend_node_index];
  blend_node.m_name = "BlendWalkRun";

  AnimNodeResource& graph_node = graph_resource.m_nodes[0];
  graph_node.m_socket_accessor->RegisterInput<AnimData>("GraphOutput", nullptr);

  REQUIRE(graph_node.m_socket_accessor->m_inputs.size() == 1);
  REQUIRE(blend_node.m_socket_accessor->GetInputIndex("Input0") == 0);
  REQUIRE(blend_node.m_socket_accessor->GetInputIndex("Input1") == 1);

  graph_resource.connectSockets(walk_node, "Output", blend_node, "Input0");
  graph_resource.connectSockets(run_node, "Output", blend_node, "Input1");
  graph_resource.connectSockets(
      blend_node,
      "Output",
      graph_resource.getGraphOutputNode(),
      "GraphOutput");

  graph_resource.saveToFile("WalkGraph.animgraph.json");
  AnimGraphResource graph_resource_loaded;
  graph_resource_loaded.loadFromFile("WalkGraph.animgraph.json");

  AnimGraph graph;
  graph_resource_loaded.createInstance(graph);
  AnimGraphContext graph_context;

  ozz::animation::Skeleton skeleton;
  REQUIRE(load_skeleton(skeleton, "data/skeleton.ozz"));
  graph_context.m_skeleton = &skeleton;

  REQUIRE(graph.init(graph_context));

  REQUIRE(graph.m_nodes.size() == 5);
  REQUIRE(graph.m_nodes[0]->m_node_type_name == "BlendTree");
  REQUIRE(graph.m_nodes[1]->m_node_type_name == "BlendTree");
  REQUIRE(graph.m_nodes[2]->m_node_type_name == "AnimSampler");
  REQUIRE(graph.m_nodes[3]->m_node_type_name == "AnimSampler");
  REQUIRE(graph.m_nodes[4]->m_node_type_name == "Blend2");

  // connections within the graph
  AnimSamplerNode* anim_sampler_walk =
      dynamic_cast<AnimSamplerNode*>(graph.m_nodes[2]);
  AnimSamplerNode* anim_sampler_run =
      dynamic_cast<AnimSamplerNode*>(graph.m_nodes[3]);
  Blend2Node* blend2_instance = dynamic_cast<Blend2Node*>(graph.m_nodes[4]);

  // check node input dependencies
  size_t anim_sampler_index0 = anim_sampler_walk->m_index;
  size_t anim_sampler_index1 = anim_sampler_run->m_index;
  size_t blend_index = blend2_instance->m_index;

  REQUIRE(graph.m_node_input_connections[blend_index].size() == 2);
  CHECK(
      graph.m_node_input_connections[blend_index][0].m_source_node
      == anim_sampler_walk);
  CHECK(
      graph.m_node_input_connections[blend_index][1].m_source_node
      == anim_sampler_run);

  REQUIRE(graph.m_node_output_connections[anim_sampler_index0].size() == 1);
  CHECK(
      graph.m_node_output_connections[anim_sampler_index0][0].m_target_node
      == blend2_instance);

  REQUIRE(graph.m_node_output_connections[anim_sampler_index1].size() == 1);
  CHECK(
      graph.m_node_output_connections[anim_sampler_index1][0].m_target_node
      == blend2_instance);

  // Ensure animation sampler nodes use the correct files
  REQUIRE(anim_sampler_walk->m_filename == "data/walk.anim.ozz");
  REQUIRE(anim_sampler_walk->m_animation != nullptr);

  REQUIRE(anim_sampler_run->m_filename == "data/run.anim.ozz");
  REQUIRE(anim_sampler_run->m_animation != nullptr);

  WHEN("Emulating Graph Evaluation") {
    CHECK(graph.m_anim_data_allocator.size() == 0);
    graph.prepareNodeEval(graph_context, walk_node_index);
    graph.finishNodeEval(walk_node_index);

    graph.prepareNodeEval(graph_context, run_node_index);
    graph.finishNodeEval(run_node_index);

    graph.prepareNodeEval(graph_context, blend_node_index);
    CHECK(blend2_instance->i_input0 == anim_sampler_walk->o_output);
    CHECK(blend2_instance->i_input1 == anim_sampler_run->o_output);

    graph.finishNodeEval(blend_node_index);

    // Evaluate output node.
    graph.evalOutputNode();
    graph.finishNodeEval(0);

    const Socket* graph_output_socket = graph.getOutputSocket("GraphOutput");
    AnimData* graph_output =
        static_cast<AnimData*>(*graph_output_socket->m_reference.ptr_ptr);

    CHECK(graph_output->m_local_matrices.size() == graph_context.m_skeleton->num_soa_joints());

    CHECK(blend2_instance->o_output == *graph_output_socket->m_reference.ptr_ptr);
  }

  graph_context.freeAnimations();
}


TEST_CASE("InputAttributeConversion", "[AnimGraphResource]") {
  int node_id = 3321;
  int input_index = 221;
  int output_index = 125;
  int parsed_node_id;
  int parsed_input_index;
  int parsed_output_index;

  int attribute_id = GenerateInputAttributeId(node_id, input_index);
  SplitInputAttributeId(attribute_id, &parsed_node_id, &parsed_input_index);
  CHECK(node_id == parsed_node_id);
  CHECK(input_index == parsed_input_index);

  attribute_id = GenerateOutputAttributeId(node_id, output_index);
  SplitOutputAttributeId(attribute_id, &parsed_node_id, &parsed_output_index);
  CHECK(node_id == parsed_node_id);
  CHECK(output_index == parsed_output_index);
}


TEST_CASE("ResourceSaveLoadMathGraphInputs", "[AnimGraphResource]") {
  AnimGraphResource graph_resource_origin;

  graph_resource_origin.clear();
  graph_resource_origin.m_name = "TestInputOutputGraph";

  size_t float_to_vec3_node_index = graph_resource_origin.addNode(
      AnimNodeResourceFactory("MathFloatToVec3Node"));

  AnimNodeResource& graph_output_node =
      graph_resource_origin.getGraphOutputNode();
  graph_output_node.m_socket_accessor->RegisterInput<float>(
      "GraphFloatOutput",
      nullptr);
  graph_output_node.m_socket_accessor->RegisterInput<Vec3>(
      "GraphVec3Output",
      nullptr);

  AnimNodeResource& graph_input_node =
      graph_resource_origin.getGraphInputNode();
  graph_input_node.m_socket_accessor->RegisterOutput<float>(
      "GraphFloatInput",
      nullptr);

  // Prepare graph inputs and outputs
  AnimNodeResource& float_to_vec3_node =
      graph_resource_origin.m_nodes[float_to_vec3_node_index];

  REQUIRE(graph_resource_origin.connectSockets(
      graph_input_node,
      "GraphFloatInput",
      graph_output_node,
      "GraphFloatOutput"));

  REQUIRE(graph_resource_origin.connectSockets(
      graph_input_node,
      "GraphFloatInput",
      float_to_vec3_node,
      "Input0"));
  REQUIRE(graph_resource_origin.connectSockets(
      graph_input_node,
      "GraphFloatInput",
      float_to_vec3_node,
      "Input1"));
  REQUIRE(graph_resource_origin.connectSockets(
      graph_input_node,
      "GraphFloatInput",
      float_to_vec3_node,
      "Input2"));

  REQUIRE(graph_resource_origin.connectSockets(
      float_to_vec3_node,
      "Output",
      graph_output_node,
      "GraphVec3Output"));

  WHEN("Saving and loading graph resource") {
    const char* filename = "ResourceSaveLoadGraphInputs.json";
    graph_resource_origin.saveToFile(filename);

    AnimGraphResource graph_resource_loaded;
    graph_resource_loaded.loadFromFile(filename);

    const AnimNodeResource& graph_loaded_output_node =
        graph_resource_loaded.m_nodes[0];
    const AnimNodeResource& graph_loaded_input_node =
        graph_resource_loaded.m_nodes[1];

    THEN("Graph inputs and outputs must be in loaded resource as well.") {
      REQUIRE(
          graph_output_node.m_socket_accessor->m_inputs.size()
          == graph_loaded_output_node.m_socket_accessor->m_inputs.size());

      REQUIRE(
          graph_input_node.m_socket_accessor->m_outputs.size()
          == graph_loaded_input_node.m_socket_accessor->m_outputs.size());

      REQUIRE(
          graph_loaded_input_node.m_socket_accessor->GetOutputSocket(
              "GraphFloatInput")
          != nullptr);

      REQUIRE(
          graph_loaded_output_node.m_socket_accessor->GetInputSocket(
              "GraphFloatOutput")
          != nullptr);
      REQUIRE(
          graph_loaded_output_node.m_socket_accessor->GetInputSocket(
              "GraphVec3Output")
          != nullptr);

      WHEN("Instantiating an AnimGraph") {
        AnimGraph anim_graph;
        graph_resource_loaded.createInstance(anim_graph);

        REQUIRE(anim_graph.getInputSocket("GraphFloatInput") != nullptr);
        REQUIRE(
            anim_graph.getInputPtr("GraphFloatInput")
            == anim_graph.m_input_buffer);

        float* graph_float_input = nullptr;
        graph_float_input =
            static_cast<float*>(anim_graph.getInputPtr("GraphFloatInput"));

        *graph_float_input = 123.456f;

        AND_WHEN("Evaluating Graph") {
          AnimGraphContext context;
          context.m_graph = &anim_graph;

          anim_graph.updateTime(0.f);
          anim_graph.evaluate(context);

          Socket* float_output_socket =
              anim_graph.getOutputSocket("GraphFloatOutput");
          Socket* vec3_output_socket =
              anim_graph.getOutputSocket("GraphVec3Output");
          Vec3& vec3_output =
              *static_cast<Vec3*>(vec3_output_socket->m_reference.ptr);

          THEN("output vector components equal the graph input vaulues") {
            CHECK(vec3_output.v[0] == *graph_float_input);
            CHECK(vec3_output.v[1] == *graph_float_input);
            CHECK(vec3_output.v[2] == *graph_float_input);
          }

          context.freeAnimations();
        }
      }
    }
  }
}

/*
TEST_CASE("SimpleMathEvaluations", "[AnimGraphResource]") {
  AnimGraphResource graph_resource_origin;

  graph_resource_origin.clear();
  graph_resource_origin.m_name = "TestInputOutputGraph";

  size_t math_add0_node_index =
      graph_resource_origin.addNode(AnimNodeResourceFactory("MathAddNode"));

  size_t math_add1_node_index =
      graph_resource_origin.addNode(AnimNodeResourceFactory("MathAddNode"));

  AnimNodeResource& graph_output_node =
      graph_resource_origin.getGraphOutputNode();

  graph_output_node.m_socket_accessor->RegisterInput<float>(
      "GraphFloat0Output",
      nullptr);
  graph_output_node.m_socket_accessor->RegisterInput<float>(
      "GraphFloat1Output",
      nullptr);
  graph_output_node.m_socket_accessor->RegisterInput<float>(
      "GraphFloat2Output",
      nullptr);

  AnimNodeResource& graph_input_node =
      graph_resource_origin.getGraphInputNode();
  graph_input_node.m_socket_accessor->RegisterOutput<float>(
      "GraphFloatInput",
      nullptr);

  // Prepare graph inputs and outputs
  AnimNodeResource& math_add0_node =
      graph_resource_origin.m_nodes[math_add0_node_index];
  AnimNodeResource& math_add1_node =
      graph_resource_origin.m_nodes[math_add1_node_index];

  // direct output
  REQUIRE(graph_resource_origin.connectSockets(
      graph_input_node,
      "GraphFloatInput",
      graph_output_node,
      "GraphFloat0Output"));

  // add0 node
  REQUIRE(graph_resource_origin.connectSockets(
      graph_input_node,
      "GraphFloatInput",
      math_add0_node,
      "Input0"));

  REQUIRE(graph_resource_origin.connectSockets(
      graph_input_node,
      "GraphFloatInput",
      math_add0_node,
      "Input1"));

  REQUIRE(graph_resource_origin.connectSockets(
      math_add0_node,
      "Output",
      graph_output_node,
      "GraphFloat1Output"));

  // add1 node
  REQUIRE(graph_resource_origin.connectSockets(
      math_add0_node,
      "Output",
      math_add1_node,
      "Input0"));

  REQUIRE(graph_resource_origin.connectSockets(
      graph_input_node,
      "GraphFloatInput",
      math_add1_node,
      "Input1"));

  REQUIRE(graph_resource_origin.connectSockets(
      math_add1_node,
      "Output",
      graph_output_node,
      "GraphFloat2Output"));

  WHEN("Saving and loading graph resource") {
    const char* filename = "ResourceSaveLoadGraphInputs.json";
    graph_resource_origin.saveToFile(filename);

    AnimGraphResource graph_resource_loaded;
    graph_resource_loaded.loadFromFile(filename);

    const AnimNodeResource& graph_loaded_output_node =
        graph_resource_loaded.m_nodes[0];
    const AnimNodeResource& graph_loaded_input_node =
        graph_resource_loaded.m_nodes[1];

    WHEN("Instantiating an AnimGraph") {
      AnimGraph anim_graph;
      graph_resource_loaded.createInstance(anim_graph);

      REQUIRE(anim_graph.getInputSocket("GraphFloatInput") != nullptr);
      REQUIRE(
          anim_graph.getInputPtr("GraphFloatInput")
          == anim_graph.m_input_buffer);

      float* graph_float_input = nullptr;
      graph_float_input =
          static_cast<float*>(anim_graph.getInputPtr("GraphFloatInput"));

      *graph_float_input = 123.456f;

      AND_WHEN("Evaluating Graph") {
        AnimGraphContext context;
        context.m_graph = &anim_graph;

        anim_graph.updateTime(0.f);
        anim_graph.evaluate(context);

        Socket* float0_output_socket =
            anim_graph.getOutputSocket("GraphFloat0Output");
        Socket* float1_output_socket =
            anim_graph.getOutputSocket("GraphFloat1Output");
        Socket* float2_output_socket =
            anim_graph.getOutputSocket("GraphFloat2Output");

        REQUIRE(float0_output_socket != nullptr);
        REQUIRE(float1_output_socket != nullptr);
        REQUIRE(float2_output_socket != nullptr);

        float& float0_output =
            *static_cast<float*>(float0_output_socket->m_reference.ptr);
        float& float1_output =
            *static_cast<float*>(float1_output_socket->m_reference.ptr);
        float& float2_output =
            *static_cast<float*>(float2_output_socket->m_reference.ptr);

        THEN("output vector components equal the graph input vaulues") {
          CHECK(float0_output == Approx(*graph_float_input));
          CHECK(float1_output == Approx(*graph_float_input * 2.));
          CHECK(float2_output == Approx(*graph_float_input * 3.));
        }

        context.freeAnimations();
      }
    }
  }
}

*/