#include <curl/curl.h>
#include <pistache/client.h>
#include <rapidjson/document.h>
#include <sys/stat.h>

#include <catch2/catch.hpp>
#include <daggy/Serialization.hpp>
#include <daggy/executors/task/ForkingTaskExecutor.hpp>
#include <daggy/executors/task/NoopTaskExecutor.hpp>
#include <daggy/loggers/dag_run/OStreamLogger.hpp>
#include <daggyd/Server.hpp>
#include <filesystem>
#include <iostream>
#include <thread>

namespace rj = rapidjson;

using namespace daggy;

TEST_CASE("rest_endpoint", "[server_basic]")
{
  std::stringstream ss;
  daggy::executors::task::ForkingTaskExecutor executor(10);
  daggy::loggers::dag_run::OStreamLogger logger(ss);
  Pistache::Address listenSpec("localhost", Pistache::Port(0));

  const size_t nDAGRunners = 10, nWebThreads = 10;

  daggy::daggyd::Server server(listenSpec, logger, executor, nDAGRunners,
                               "/dev/null");
  server.init(nWebThreads);
  server.start();

  const std::string host    = "localhost:";
  const std::string baseURL = host + std::to_string(server.getPort());

  SECTION("Ready Endpoint")
  {
    auto response = HTTP_REQUEST(baseURL + "/ready");
    REQUIRE(response.code == HTTPCode::Ok);
  }

  SECTION("Querying a non-existent dagrunid should fail ")
  {
    auto response = HTTP_REQUEST(baseURL + "/v1/dagrun/100");
    REQUIRE(response.code != HTTPCode::Ok);
  }

  SECTION("Simple DAGRun Submission")
  {
    std::string dagRun = R"({
            "tag": "unit_server",
            "parameters": { "FILE": [ "A", "B" ] },
            "tasks": {
               "touch": { "job": { "command": [ "/bin/touch", "dagrun_{{FILE}}" ], "environment": []} },
               "cat": {  "job": { "command": [ "/bin/cat", "dagrun_A", "dagrun_B" ], "environment": []},
                 "parents": [ "touch" ]
               }
               }
        })";

    auto dagSpec = daggy::dagFromJSON(dagRun);

    // Submit, and get the runID
    daggy::DAGRunID runID = 0;
    {
      auto response = HTTP_REQUEST(baseURL + "/v1/dagrun/", dagRun, "POST");
      REQUIRE(response.code == HTTPCode::Ok);

      rj::Document doc;
      daggy::checkRJParse(doc.Parse(response.body.c_str()));
      REQUIRE(doc.IsObject());
      REQUIRE(doc.HasMember("runID"));

      runID = doc["runID"].GetUint64();
    }

    // Ensure our runID shows up in the list of running DAGs
    {
      auto response = HTTP_REQUEST(baseURL + "/v1/dagruns?all=1");
      REQUIRE(response.code == HTTPCode::Ok);

      rj::Document doc;
      daggy::checkRJParse(doc.Parse(response.body.c_str()));
      REQUIRE(doc.IsArray());
      REQUIRE(doc.Size() >= 1);

      // Ensure that our DAG is in the list and matches our given DAGRunID
      bool found       = false;
      const auto &runs = doc.GetArray();
      for (size_t i = 0; i < runs.Size(); ++i) {
        const auto &run = runs[i];
        REQUIRE(run.IsObject());
        REQUIRE(run.HasMember("tag"));
        REQUIRE(run.HasMember("runID"));

        std::string runName = run["tag"].GetString();
        if (runName == "unit_server") {
          REQUIRE(run["runID"].GetUint64() == runID);
          found = true;
          break;
        }
      }
      REQUIRE(found);
    }

    // Ensure we can get one of our tasks
    {
      auto response = HTTP_REQUEST(baseURL + "/v1/dagrun/" +
                                   std::to_string(runID) + "/task/cat_0");
      REQUIRE(response.code == HTTPCode::Ok);

      rj::Document doc;
      daggy::checkRJParse(doc.Parse(response.body.c_str()));

      REQUIRE_NOTHROW(daggy::taskFromJSON("cat", doc));
      auto task = daggy::taskFromJSON("cat", doc);

      REQUIRE(task == dagSpec.tasks.at("cat"));
    }

    // Wait until our DAG is complete
    bool complete = true;
    for (auto i = 0; i < 10; ++i) {
      auto response =
          HTTP_REQUEST(baseURL + "/v1/dagrun/" + std::to_string(runID));
      REQUIRE(response.code == HTTPCode::Ok);
      rj::Document doc;
      daggy::checkRJParse(doc.Parse(response.body.c_str()));
      REQUIRE(doc.IsObject());

      REQUIRE(doc.HasMember("taskStates"));
      const auto &taskStates = doc["taskStates"].GetObject();

      size_t nStates = 0;
      for (auto it = taskStates.MemberBegin(); it != taskStates.MemberEnd();
           ++it) {
        nStates++;
      }
      REQUIRE(nStates == 3);

      complete = true;
      for (auto it = taskStates.MemberBegin(); it != taskStates.MemberEnd();
           ++it) {
        std::string state = it->value.GetString();
        if (state != "COMPLETED") {
          complete = false;
          break;
        }
      }
      if (complete)
        break;
      std::this_thread::sleep_for(std::chrono::seconds(1));
    }
    REQUIRE(complete);

    std::this_thread::sleep_for(std::chrono::seconds(2));
    for (const auto &pth : std::vector<fs::path>{"dagrun_A", "dagrun_B"}) {
      REQUIRE(fs::exists(pth));
      fs::remove(pth);
    }
  }
}

TEST_CASE("Server cancels and resumes execution", "[server_resume]")
{
  std::stringstream ss;
  daggy::executors::task::ForkingTaskExecutor executor(10);
  daggy::loggers::dag_run::OStreamLogger logger(std::cout);
  Pistache::Address listenSpec("localhost", Pistache::Port(0));

  const size_t nDAGRunners = 10, nWebThreads = 10;

  daggy::daggyd::Server server(listenSpec, logger, executor, nDAGRunners,
                               "/dev/null");
  server.init(nWebThreads);
  server.start();

  const std::string host    = "localhost:";
  const std::string baseURL = host + std::to_string(server.getPort());

  SECTION("Cancel / Resume DAGRun")
  {
    std::string dagRunJSON = R"({
            "tag": "unit_server",
            "tasks": {
               "touch_A": { "job": { "command": [ "/bin/touch", "resume_touch_a" ]}, "children": ["touch_C"] },
               "sleep_B": { "job": { "command": [ "/bin/sleep", "3" ]}, "children": ["touch_C"] },
               "touch_C": { "job": { "command": [ "/bin/touch", "resume_touch_c" ]} }
               }
        })";

    auto dagSpec = daggy::dagFromJSON(dagRunJSON);

    // Submit, and get the runID
    daggy::DAGRunID runID;
    {
      auto response = HTTP_REQUEST(baseURL + "/v1/dagrun/", dagRunJSON, "POST");
      REQUIRE(response.code == HTTPCode::Ok);

      rj::Document doc;
      daggy::checkRJParse(doc.Parse(response.body.c_str()));
      REQUIRE(doc.IsObject());
      REQUIRE(doc.HasMember("runID"));

      runID = doc["runID"].GetUint64();
    }

    std::this_thread::sleep_for(1s);

    // Stop the current run
    {
      auto response = HTTP_REQUEST(
          baseURL + "/v1/dagrun/" + std::to_string(runID) + "/state/KILLED", "",
          "PATCH");
      REQUIRE(response.code == HTTPCode::Ok);
      REQUIRE(logger.getDAGRunState(runID) == +daggy::RunState::KILLED);
    }

    // Verify that the run still exists
    {
      auto dagRun = logger.getDAGRun(runID);
      REQUIRE(dagRun.taskRunStates.at("touch_A_0") ==
              +daggy::RunState::COMPLETED);
      REQUIRE(fs::exists("resume_touch_a"));

      REQUIRE(dagRun.taskRunStates.at("sleep_B_0") ==
              +daggy::RunState::ERRORED);
      REQUIRE(dagRun.taskRunStates.at("touch_C_0") == +daggy::RunState::QUEUED);
    }

    // Set the errored task state
    {
      auto url = baseURL + "/v1/dagrun/" + std::to_string(runID) +
                 "/task/sleep_B_0/state/QUEUED";
      auto response = HTTP_REQUEST(url, "", "PATCH");
      REQUIRE(response.code == HTTPCode::Ok);
      REQUIRE(logger.getTaskState(runID, "sleep_B_0") ==
              +daggy::RunState::QUEUED);
    }

    // Resume
    {
      struct stat s;

      lstat("resume_touch_A", &s);
      auto preMTime = s.st_mtim.tv_sec;

      auto response = HTTP_REQUEST(
          baseURL + "/v1/dagrun/" + std::to_string(runID) + "/state/QUEUED", "",
          "PATCH");

      // Wait for run to complete
      std::this_thread::sleep_for(3s);
      REQUIRE(logger.getDAGRunState(runID) == +daggy::RunState::COMPLETED);

      REQUIRE(fs::exists("resume_touch_c"));
      REQUIRE(fs::exists("resume_touch_a"));

      for (const auto &[taskName, task] : dagSpec.tasks) {
        REQUIRE(logger.getTaskState(runID, taskName + "_0") ==
                +daggy::RunState::COMPLETED);
      }

      // Ensure "touch_A" wasn't run again
      lstat("resume_touch_A", &s);
      auto postMTime = s.st_mtim.tv_sec;
      REQUIRE(preMTime == postMTime);
    }
  }

  server.shutdown();
}