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

#include "AnimGraphResource.h"
#include "catch.hpp"

using namespace AniGraph;

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";
  AnimNodeResource& run_node = graph_resource.m_nodes[run_node_index];
  walk_node.m_name = "RunAnim";
  AnimNodeResource& blend_node = graph_resource.m_nodes[blend_node_index];
  walk_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);

  AnimGraphConnection walk_to_blend;
  walk_to_blend.m_source_node = &walk_node;
  walk_to_blend.m_source_socket =
      walk_node.m_socket_accessor->FindOutputSocket("Output");
  walk_to_blend.m_target_node = &blend_node;
  walk_to_blend.m_target_socket =
      blend_node.m_socket_accessor->FindInputSocket("Input0");
  graph_resource.m_connections.push_back(walk_to_blend);

  AnimGraphConnection run_to_blend;
  run_to_blend.m_source_node = &run_node;
  run_to_blend.m_source_socket =
      run_node.m_socket_accessor->FindOutputSocket("Output");
  run_to_blend.m_target_node = &blend_node;
  run_to_blend.m_target_socket =
      blend_node.m_socket_accessor->FindInputSocket("Input1");
  graph_resource.m_connections.push_back(run_to_blend);

  AnimGraphConnection blend_to_output;
  blend_to_output.m_source_node = &blend_node;
  blend_to_output.m_source_socket =
      blend_node.m_socket_accessor->FindOutputSocket("Output");
  blend_to_output.m_target_node = &graph_resource.m_nodes[0];
  blend_to_output.m_target_socket =
      graph_node.m_socket_accessor->FindInputSocket("GraphOutput");
  graph_resource.m_connections.push_back(blend_to_output);

  graph_resource.saveToFile("WalkGraph.animgraph.json");

  AnimGraph graph = AnimGraph::createFromResource(graph_resource);

  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_instance0 =
      dynamic_cast<AnimSamplerNode*>(graph.m_nodes[2]);
  AnimSamplerNode* anim_sampler_instance1 =
      dynamic_cast<AnimSamplerNode*>(graph.m_nodes[3]);
  Blend2Node* blend2_instance = dynamic_cast<Blend2Node*>(graph.m_nodes[4]);
  CHECK(anim_sampler_instance0->m_output == &blend2_instance->m_input0);
  CHECK(anim_sampler_instance1->m_output == &blend2_instance->m_input1);

  // connections from graph to the graph node
  CHECK(graph.m_socket_accessor->m_inputs.size() == 1);
  CHECK(graph.m_socket_accessor->FindInputSocket("GraphOutput"));
  CHECK(
      reinterpret_cast<char*>(blend2_instance->m_output)
      == graph.getOutput("GraphOutput"));

  // check node input dependencies
  size_t anim_sampler_index0 = anim_sampler_instance0->m_index;
  size_t anim_sampler_index1 = anim_sampler_instance1->m_index;
  size_t blend_index = blend2_instance->m_index;

  CHECK(graph.m_node_inputs[anim_sampler_index0].size() == 0);
  CHECK(graph.m_node_inputs[anim_sampler_index1].size() == 0);
  CHECK(graph.m_node_inputs[blend_index].size() == 3);
  CHECK(graph.m_node_inputs[blend_index][0].m_node == anim_sampler_instance0);
  CHECK(graph.m_node_inputs[blend_index][1].m_node == anim_sampler_instance1);
  CHECK(graph.m_node_inputs[blend_index][2].m_node == nullptr);
}

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("ResourceSaveLoadGraphInputs", "[AnimGraphResource]") {
  AnimGraphResource graph_resource_origin;

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

  AnimNodeResource& graph_output_node = graph_resource_origin.m_nodes[0];
  graph_output_node.m_socket_accessor->RegisterInput<AnimData>(
      "GraphOutput",
      nullptr);
  graph_output_node.m_socket_accessor->RegisterInput<float>(
      "SomeFloatOutput",
      nullptr);

  AnimNodeResource& graph_input_node = graph_resource_origin.m_nodes[1];
  graph_input_node.m_socket_accessor->RegisterOutput<AnimData>(
      "GraphAnimInput",
      nullptr);
  graph_input_node.m_socket_accessor->RegisterOutput<float>(
      "GraphFloatInput",
      nullptr);
  graph_input_node.m_socket_accessor->RegisterOutput<bool>(
      "GraphBoolInput",
      nullptr);

  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->FindOutputSocket(
              "GraphAnimInput")
          != nullptr);
      REQUIRE(
          graph_loaded_input_node.m_socket_accessor->FindOutputSocket(
              "GraphFloatInput")
          != nullptr);
      REQUIRE(
          graph_loaded_input_node.m_socket_accessor->FindOutputSocket(
              "GraphBoolInput")
          != nullptr);

      REQUIRE(
          graph_loaded_output_node.m_socket_accessor->FindInputSocket(
              "GraphOutput")
          != nullptr);
      REQUIRE(
          graph_loaded_output_node.m_socket_accessor->FindInputSocket(
              "SomeFloatOutput")
          != nullptr);
    }

    WHEN("Instantiating an AnimGraph") {
      AnimGraph anim_graph =
          AnimGraph::createFromResource(graph_resource_loaded);
      REQUIRE(
          anim_graph.getOutput("GraphOutput") == anim_graph.m_output_buffer);
      REQUIRE(
          anim_graph.getOutput("SomeFloatOutput")
          == anim_graph.m_output_buffer + sizeof(AnimData));

      REQUIRE(anim_graph.getInput("GraphAnimInput") == nullptr);
      REQUIRE(anim_graph.getInput("GraphFloatInput") == nullptr);
      REQUIRE(anim_graph.getInput("GraphBoolInput") == nullptr);
    }
  }

  WHEN("Connecting input to output and instantiating the graph") {
    AnimNodeResource& graph_output_node = graph_resource_origin.m_nodes[0];
    AnimNodeResource& graph_input_node = graph_resource_origin.m_nodes[1];

    REQUIRE(graph_resource_origin.connectSockets(
        graph_input_node,
        "GraphAnimInput",
        graph_output_node,
        "GraphOutput"));

    AnimGraph anim_graph = AnimGraph::createFromResource(graph_resource_origin);

    void* graph_anim_input_ptr = anim_graph.getInput("GraphAnimInput");
    void* graph_output_ptr = anim_graph.getOutput("GraphOutput");

    REQUIRE(graph_anim_input_ptr == graph_output_ptr);
    REQUIRE(graph_output_ptr == anim_graph.m_output_buffer);

    REQUIRE(
        anim_graph.getInput("GraphAnimInput")
        == anim_graph.getOutput("GraphOutput"));
  }
}

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

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

  AnimNodeResource& graph_output_node = graph_resource.m_nodes[0];
  graph_output_node.m_socket_accessor->RegisterInput<float>(
      "GraphFloatOutput",
      nullptr);
  graph_output_node.m_socket_accessor->RegisterInput<AnimData>(
      "GraphAnimOutput",
      nullptr);

  AnimNodeResource& graph_input_node = graph_resource.m_nodes[1];
  graph_input_node.m_socket_accessor->RegisterOutput<float>(
      "GraphFloatInput",
      nullptr);
  graph_input_node.m_socket_accessor->RegisterOutput<float>(
      "SpeedScaleInput",
      nullptr);
  graph_input_node.m_socket_accessor->RegisterOutput<AnimData>(
      "GraphAnimInput0",
      nullptr);
  graph_input_node.m_socket_accessor->RegisterOutput<AnimData>(
      "GraphAnimInput1",
      nullptr);

  WHEN("Connecting float input with float output") {
    REQUIRE(graph_resource.connectSockets(
        graph_resource.getGraphInputNode(),
        "GraphFloatInput",
        graph_resource.getGraphOutputNode(),
        "GraphFloatOutput"));

    AnimGraph anim_graph = AnimGraph::createFromResource(graph_resource);

    THEN("Writing to the input pointer changes the value of the output.") {
      float* float_input_ptr = (float*)anim_graph.getInput("GraphFloatInput");
      REQUIRE(float_input_ptr != nullptr);
      *float_input_ptr = 23.123f;

      float* float_output_ptr =
          (float*)anim_graph.getOutput("GraphFloatOutput");
      REQUIRE(float_output_ptr != nullptr);
      CHECK(*float_output_ptr == Approx(23.123f));
    }
  }

  WHEN("Connecting adding a Blend2 node") {
    size_t blend2_node_index =
        graph_resource.addNode(AnimNodeResourceFactory("Blend2"));
    AnimNodeResource& blend2_node_resource =
        graph_resource.m_nodes[blend2_node_index];

    REQUIRE(graph_resource.connectSockets(
        graph_resource.getGraphInputNode(),
        "GraphFloatInput",
        blend2_node_resource,
        "Weight"));

    THEN("Connected float input points to the blend weight.") {
      AnimGraph anim_graph = AnimGraph::createFromResource(graph_resource);
      Blend2Node* blend2_node =
          dynamic_cast<Blend2Node*>(anim_graph.m_nodes[blend2_node_index]);

      REQUIRE(
          *anim_graph.m_socket_accessor->m_outputs[0].m_value.ptr_ptr
          == &blend2_node->m_blend_weight);
      float* float_input_ptr = (float*)anim_graph.getInput("GraphFloatInput");
      REQUIRE(float_input_ptr == &blend2_node->m_blend_weight);
    }

    WHEN(
        "Connecting AnimData inputs to blend2 node and blend2 output to graph "
        "output.") {
      REQUIRE(graph_resource.connectSockets(
          graph_resource.getGraphInputNode(),
          "GraphAnimInput0",
          blend2_node_resource,
          "Input0"));

      REQUIRE(graph_resource.connectSockets(
          graph_resource.getGraphInputNode(),
          "GraphAnimInput1",
          blend2_node_resource,
          "Input1"));

      REQUIRE(graph_resource.connectSockets(
          blend2_node_resource,
          "Output",
          graph_resource.getGraphOutputNode(),
          "GraphAnimOutput"));

      THEN(
          "AnimData from output gets blended and result is written to "
          "Output.") {
        AnimGraph anim_graph = AnimGraph::createFromResource(graph_resource);
        Blend2Node* blend2_node =
            dynamic_cast<Blend2Node*>(anim_graph.m_nodes[blend2_node_index]);

        AnimData* graph_input0 =
            (AnimData*)anim_graph.getInput("GraphAnimInput0");
        REQUIRE(graph_input0 == &blend2_node->m_input0);
        REQUIRE(
            anim_graph.m_nodes[1]
            == anim_graph.getAnimNodeForInput(blend2_node_index, "Input0"));

        AnimData* graph_input1 =
            (AnimData*)anim_graph.getInput("GraphAnimInput1");
        REQUIRE(graph_input1 == &blend2_node->m_input1);
        REQUIRE(
            anim_graph.m_nodes[1]
            == anim_graph.getAnimNodeForInput(blend2_node_index, "Input1"));

        AnimData* graph_output =
            (AnimData*)anim_graph.getOutput("GraphAnimOutput");
        REQUIRE(graph_output == blend2_node->m_output);
        REQUIRE(
            anim_graph.m_nodes[blend2_node_index]
            == anim_graph.getAnimNodeForInput(0, "GraphAnimOutput"));
      }
    }
  }

  WHEN("Adding AnimSampler Nodes") {
    size_t blend2_node_index =
        graph_resource.addNode(AnimNodeResourceFactory("Blend2"));
    size_t sampler_node_index =
        graph_resource.addNode(AnimNodeResourceFactory("AnimSampler"));
    size_t speed_scale_node_index =
        graph_resource.addNode(AnimNodeResourceFactory("SpeedScale"));

    AnimNodeResource& blend2_node_resource =
        graph_resource.m_nodes[blend2_node_index];
    AnimNodeResource& sampler_node_resource =
        graph_resource.m_nodes[sampler_node_index];
    AnimNodeResource& speed_scale_node_resource =
        graph_resource.m_nodes[speed_scale_node_index];

    REQUIRE(graph_resource.connectSockets(
        graph_resource.getGraphInputNode(),
        "GraphFloatInput",
        blend2_node_resource,
        "Weight"));

    REQUIRE(graph_resource.connectSockets(
        graph_resource.getGraphInputNode(),
        "SpeedScaleInput",
        speed_scale_node_resource,
        "SpeedScale"));

    REQUIRE(graph_resource.connectSockets(
        graph_resource.getGraphInputNode(),
        "GraphAnimInput0",
        blend2_node_resource,
        "Input0"));

    REQUIRE(graph_resource.connectSockets(
        sampler_node_resource,
        "Output",
        speed_scale_node_resource,
        "Input"));

    REQUIRE(graph_resource.connectSockets(
        speed_scale_node_resource,
        "Output",
        blend2_node_resource,
        "Input1"));

    REQUIRE(graph_resource.connectSockets(
        blend2_node_resource,
        "Output",
        graph_resource.getGraphOutputNode(),
        "GraphAnimOutput"));

    THEN("Data flow and node ordering must be correct.") {
      AnimGraph anim_graph = AnimGraph::createFromResource(graph_resource);
      Blend2Node* blend2_node =
          dynamic_cast<Blend2Node*>(anim_graph.m_nodes[blend2_node_index]);
      SpeedScaleNode* speed_scale_node = dynamic_cast<SpeedScaleNode*>(
          anim_graph.m_nodes[speed_scale_node_index]);
      AnimSamplerNode* sampler_node = dynamic_cast<AnimSamplerNode*>(
          anim_graph.m_nodes[sampler_node_index]);

      //
      // check connectivity
      //
      AnimData* graph_input0 =
          (AnimData*)anim_graph.getInput("GraphAnimInput0");
      REQUIRE(graph_input0 == &blend2_node->m_input0);
      REQUIRE(
          anim_graph.m_nodes[1]
          == anim_graph.getAnimNodeForInput(blend2_node_index, "Input0"));

      AnimData* graph_input1 =
          (AnimData*)anim_graph.getInput("GraphAnimInput1");
      REQUIRE(graph_input1 == nullptr);

      REQUIRE(sampler_node->m_output == &speed_scale_node->m_input);
      REQUIRE(
          sampler_node
          == anim_graph.getAnimNodeForInput(speed_scale_node_index, "Input"));

      REQUIRE(speed_scale_node->m_output == &blend2_node->m_input1);
      REQUIRE(
          speed_scale_node
          == anim_graph.getAnimNodeForInput(blend2_node_index, "Input1"));

      AnimData* graph_output =
          (AnimData*)anim_graph.getOutput("GraphAnimOutput");
      REQUIRE(graph_output == blend2_node->m_output);
      REQUIRE(
          anim_graph.m_nodes[blend2_node_index]
          == anim_graph.getAnimNodeForInput(0, "GraphAnimOutput"));

      //
      // check ordering
      //
      REQUIRE(
          anim_graph.getNodeEvalOrderIndex(blend2_node)
          < anim_graph.getNodeEvalOrderIndex(sampler_node));
      REQUIRE(
          anim_graph.getNodeEvalOrderIndex(blend2_node)
          < anim_graph.getNodeEvalOrderIndex(speed_scale_node));
      REQUIRE(
          anim_graph.getNodeEvalOrderIndex(speed_scale_node)
          < anim_graph.getNodeEvalOrderIndex(sampler_node));
    }

    WHEN("Instantiating graph") {
      AnimGraph anim_graph = AnimGraph::createFromResource(graph_resource);
      float* blend_weight_input =
          reinterpret_cast<float*>(anim_graph.getInput("GraphFloatInput"));

      Blend2Node* blend2_node =
          dynamic_cast<Blend2Node*>(anim_graph.m_nodes[blend2_node_index]);
      SpeedScaleNode* speed_scale_node = dynamic_cast<SpeedScaleNode*>(
          anim_graph.m_nodes[speed_scale_node_index]);
      AnimSamplerNode* sampler_node = dynamic_cast<AnimSamplerNode*>(
          anim_graph.m_nodes[sampler_node_index]);

      WHEN("Setting weight to 0. and marking nodes active.") {
        *blend_weight_input = 0.;
        anim_graph.markActiveNodes();

        THEN("Speed scale and sampler node are inactive") {
          REQUIRE(anim_graph.checkIsNodeActive(speed_scale_node) == false);
          REQUIRE(anim_graph.checkIsNodeActive(sampler_node) == false);
        }
      }

      WHEN("Setting weight to 0. and marking nodes active") {
        *blend_weight_input = 0.1;
        anim_graph.markActiveNodes();

        THEN("Speed scale and sampler nodes are active") {
          REQUIRE(anim_graph.checkIsNodeActive(speed_scale_node) == true);
          REQUIRE(anim_graph.checkIsNodeActive(sampler_node) == true);
        }
      }

      WHEN("Setting weight to 1. and marking nodes active") {
        *blend_weight_input = 1.0;
        anim_graph.markActiveNodes();

        THEN("Speed scale and sampler nodes are active") {
          REQUIRE(anim_graph.checkIsNodeActive(speed_scale_node) == true);
          REQUIRE(anim_graph.checkIsNodeActive(sampler_node) == true);
        }
      }

      WHEN("Updating time with dt = 0.3f and speed scale = 1.0f") {
        float* speed_scale_input =
            reinterpret_cast<float*>(anim_graph.getInput("SpeedScaleInput"));

        *blend_weight_input = 0.1;
        *speed_scale_input = 1.0f;

        anim_graph.markActiveNodes();
        anim_graph.updateTime(0.3f);

        THEN ("Anim sampler node time now must be 0.3f") {
          REQUIRE(sampler_node->m_time_now == Approx(0.3f));
        }
      }

      WHEN("Updating time with dt = 0.3f and speed scale = 1.3f") {
        float* speed_scale_input =
            reinterpret_cast<float*>(anim_graph.getInput("SpeedScaleInput"));

        *blend_weight_input = 0.1;
        *speed_scale_input = 1.3f;

        anim_graph.markActiveNodes();
        anim_graph.updateTime(0.3f);

        THEN ("Anim sampler node time now must be 0.39f") {
          REQUIRE(sampler_node->m_time_now == Approx(0.39f));
        }
      }
    }
  }
}