#include <catch.hpp>
#include <atomic>

#include "pls/pls.h"

using namespace pls::internal::scheduling;
using namespace pls::internal::scheduling::lock_free;

constexpr int MAX_NUM_TASKS = 32;
constexpr int MAX_STACK_SIZE = 1024 * 8;

TEST_CASE("scheduler correctly initializes", "[internal/scheduling/scheduler]") {
  const unsigned num_tasks = 16;
  const unsigned num_threads = 4;

  pls::scheduler scheduler{num_threads, num_tasks, 4096, false};

  SECTION("task chains are valid") {
    for (unsigned i = 0; i < num_threads; i++) {
      task_manager &manager = scheduler.task_manager_for(i);
      for (unsigned j = 0; j < num_tasks; j++) {
        REQUIRE(manager.get_task(j)->depth_ == j);
        REQUIRE(manager.get_task(j)->thread_id_ == i);

        if (j < num_tasks - 1) {
          REQUIRE(manager.get_task(j)->next_ == manager.get_task(j + 1));
        }
        if (j > 0) {
          REQUIRE(manager.get_task(j)->prev_ == manager.get_task(j - 1));
        }
      }
    }
  }
}

TEST_CASE("tasks distributed over workers (do not block)", "[internal/scheduling/scheduler.h]") {
  scheduler scheduler{3, MAX_NUM_TASKS, MAX_STACK_SIZE};

  std::atomic<int> num_run{0};
  scheduler.perform_work([&] {
    scheduler::spawn([&] {
      num_run++;
      while (num_run < 3);
    });
    scheduler::spawn([&] {
      while (num_run < 1);
      num_run++;
      while (num_run < 3);
    });
    scheduler::spawn([&] {
      while (num_run < 2);
      num_run++;
    });
    scheduler::sync();
  });
  REQUIRE(num_run == 3);
}

unsigned fib_serial(unsigned n) {
  if (n <= 1) {
    return n;
  }

  return fib_serial(n - 1) + fib_serial(n - 2);
}

unsigned fib_pls(unsigned n) {
  if (n <= 1) {
    return n;
  }

  unsigned a, b;
  pls::invoke(
      [&a, n] { a = fib_pls(n - 1); },
      [&b, n] { b = fib_pls(n - 2); }
  );
  return a + b;
}

TEST_CASE("simple fib", "[internal/scheduling/scheduler]") {
  scheduler scheduler{3, MAX_NUM_TASKS, MAX_STACK_SIZE};

  scheduler.perform_work([&] {
    REQUIRE(fib_serial(28) == fib_pls(28));
  });
}