#include <catch2/catch.hpp>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <sstream>

#include "daggy/Serialization.hpp"
#include "daggy/loggers/dag_run/OStreamLogger.hpp"
#include "daggy/loggers/dag_run/RedisLogger.hpp"

using namespace daggy;
using namespace daggy::loggers::dag_run;

const TaskSet SAMPLE_TASKS{
    {"work_a",
     Task{.definedName{"work_a"},
          .job{{"command", std::vector<std::string>{"/bin/echo", "a"}}},
          .children{"c"}}},
    {"work_b",
     Task{.definedName{"work_b"},
          .job{{"command", std::vector<std::string>{"/bin/echo", "b"}}},
          .children{"c"}}},
    {"work_c",
     Task{.definedName{"work_c"},
          .job{{"command", std::vector<std::string>{"/bin/echo", "c"}}}}}};

namespace {
  void testDAGRunInit(DAGRunLogger &logger, const std::string &tag,
                      const TaskSet &tasks)
  {
    auto runID = logger.startDAGRun(DAGSpec{.tag = tag, .tasks = tasks});

    // Verify run shows up in the list
    SECTION("New run shows up in list of runs")
    {
      auto runs = logger.queryDAGRuns();
      REQUIRE(!runs.empty());
      auto it = std::find_if(runs.begin(), runs.end(), [runID](const auto &r) {
        return r.runID == runID;
      });
      REQUIRE(it != runs.end());
      REQUIRE(it->tag == tag);
      REQUIRE(it->runState == +RunState::QUEUED);
    }

    // Verify dagSpec matches
    SECTION("Can retrieve DAG Spec")
    {
      auto spec = logger.getDAGSpec(runID);
      REQUIRE(spec.tag == tag);
      REQUIRE(spec.tasks == tasks);
    }

    // Verify states
    SECTION("DAG State matches expectations")
    {
      REQUIRE(logger.getDAGRunState(runID) == +RunState::QUEUED);
      for (const auto &[k, _] : tasks) {
        REQUIRE(logger.getTaskState(runID, k) == +RunState::QUEUED);
      }
    }

    // Verify integrity of run
    SECTION("Can retrieve the full run")
    {
      auto dagRun = logger.getDAGRun(runID);

      REQUIRE(dagRun.dagSpec.tag == tag);
      REQUIRE(dagRun.dagSpec.tasks == tasks);

      REQUIRE(dagRun.taskRunStates.size() == tasks.size());
      auto nonQueuedTask = std::find_if(
          dagRun.taskRunStates.begin(), dagRun.taskRunStates.end(),
          [](const auto &a) { return a.second != +RunState::QUEUED; });

      REQUIRE(nonQueuedTask == dagRun.taskRunStates.end());
      REQUIRE(dagRun.dagStateChanges.size() == 1);
      REQUIRE(dagRun.dagStateChanges.back().state == +RunState::QUEUED);
    }

    // Update DAG state and ensure that it's updated;
    SECTION("Can update DAG state and retrieve new state")
    {
      logger.updateDAGRunState(runID, RunState::RUNNING);
      auto dagRun = logger.getDAGRun(runID);
      REQUIRE(dagRun.dagStateChanges.back().state == +RunState::RUNNING);
    }

    // Update a task state
    SECTION("Can update task state and retrieve new state")
    {
      for (const auto &[k, v] : tasks)
        logger.updateTaskState(runID, k, RunState::RUNNING);
      auto dagRun = logger.getDAGRun(runID);
      for (const auto &[k, v] : tasks) {
        REQUIRE(dagRun.taskRunStates.at(k) == +RunState::RUNNING);
      }
    }

    SECTION("Log task attempt and retrieve it")
    {
      std::string error = "long error string\nwith new\n lines";
      logger.logTaskAttempt(runID, "work_a",
                            AttemptRecord{.rc = 2, .errorLog = error});
      auto dagRun = logger.getDAGRun(runID);

      REQUIRE(dagRun.taskAttempts["work_a"].size() == 1);
      REQUIRE(dagRun.taskAttempts["work_a"][0].errorLog == error);
      REQUIRE(dagRun.taskAttempts["work_a"][0].rc == 2);
    }
  }
}  // namespace

TEST_CASE("ostream_logger", "[ostream_logger]")
{
  std::stringstream ss;
  daggy::loggers::dag_run::OStreamLogger logger(ss);

  testDAGRunInit(logger, "init_test", SAMPLE_TASKS);
}

#ifdef DAGGY_ENABLE_REDIS
TEST_CASE("redis_logger", "[redis_logger]")
{
  daggy::loggers::dag_run::RedisLogger logger;

  testDAGRunInit(logger, "init_test", SAMPLE_TASKS);
}
#endif