#include <catch.hpp>

#include <atomic>

#include "pls/pls.h"
#include "pls/internal/scheduling/task_manager.h"
#include "pls/internal/scheduling/thread_state.h"

#if PLS_DEQUE_VARIANT == PLS_DEQUE_LOCK_FREE
#include "pls/internal/scheduling/lock_free/traded_cas_field.h"
#include "pls/internal/scheduling/lock_free/external_trading_deque.h"

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

TEST_CASE("traded cas field bitmaps correctly", "[internal/scheduling/lock_free/traded_cas_field]") {
  traded_cas_field empty_field;
  REQUIRE(empty_field.is_empty());
  REQUIRE(!empty_field.is_filled_with_trade_request());
  REQUIRE(!empty_field.is_filled_with_object());

  const int stamp = 42;
  const int ID = 10;
  traded_cas_field tag_field;
  tag_field.fill_with_trade_request(stamp, ID);
  REQUIRE(tag_field.is_filled_with_trade_request());
  REQUIRE(!tag_field.is_empty());
  REQUIRE(!tag_field.is_filled_with_object());
  REQUIRE(tag_field.get_stamp() == stamp);
  REQUIRE(tag_field.get_trade_request_thread_id() == ID);

  alignas(64) task obj{nullptr, 0, 0, 0};
  traded_cas_field obj_field = tag_field;
  obj_field.fill_with_task(obj.thread_id_);
  REQUIRE(obj_field.is_filled_with_object());
  REQUIRE(!obj_field.is_empty());
  REQUIRE(!obj_field.is_filled_with_trade_request());
}

TEST_CASE("task resource stack", "[internal/scheduling/lock_free/task]") {
  // simulate scheduler with four threads and depth 1. We are thread 0.
  pls::scheduler scheduler{4, 1, 4096, false};
  pls::internal::scheduling::thread_state::set(&scheduler.thread_state_for(0));

  task *tasks[] = {scheduler.task_manager_for(0).get_task(0),
                   scheduler.task_manager_for(1).get_task(0),
                   scheduler.task_manager_for(2).get_task(0),
                   scheduler.task_manager_for(3).get_task(0)};

  SECTION("simple push/pop") {
    tasks[1]->prepare_for_push(0);
    tasks[0]->push_task_chain(tasks[1], 0);

    REQUIRE(tasks[0]->pop_task_chain() == tasks[1]);
    REQUIRE(tasks[0]->pop_task_chain() == nullptr);
  }

  SECTION("empty pop and multi push") {
    tasks[1]->prepare_for_push(0);
    tasks[0]->push_task_chain(tasks[1], 0);
    tasks[2]->prepare_for_push(0);
    tasks[0]->push_task_chain(tasks[2], 0);

    REQUIRE(tasks[0]->pop_task_chain() == tasks[2]);
    REQUIRE(tasks[0]->pop_task_chain() == tasks[1]);

    tasks[1]->prepare_for_push(0);
    tasks[0]->push_task_chain(tasks[1], 0);

    REQUIRE(tasks[0]->pop_task_chain() == tasks[1]);
    REQUIRE(tasks[0]->pop_task_chain() == nullptr);
  }

  SECTION("multiple pushes") {
    tasks[1]->prepare_for_push(0);
    tasks[0]->push_task_chain(tasks[1], 0);
    tasks[2]->prepare_for_push(0);
    tasks[0]->push_task_chain(tasks[2], 0);
    tasks[3]->prepare_for_push(0);
    tasks[0]->push_task_chain(tasks[3], 0);

    REQUIRE(tasks[0]->pop_task_chain() == tasks[3]);
    REQUIRE(tasks[0]->pop_task_chain() == tasks[2]);
    REQUIRE(tasks[0]->pop_task_chain() == tasks[1]);
    REQUIRE(tasks[0]->pop_task_chain() == nullptr);
  }
}

TEST_CASE("external trading deque", "[internal/scheduling/lock_free/external_trading_deque]") {
  // simulate scheduler with four threads and depth 1. We are thread 0.
  pls::scheduler scheduler{4, 4, 4096, false};
  pls::internal::scheduling::thread_state::set(&scheduler.thread_state_for(0));

  task *tasks_1[] = {scheduler.task_manager_for(0).get_task(0),
                     scheduler.task_manager_for(0).get_task(1),
                     scheduler.task_manager_for(0).get_task(2),
                     scheduler.task_manager_for(0).get_task(3)};
  task *tasks_2[] = {scheduler.task_manager_for(1).get_task(0),
                     scheduler.task_manager_for(1).get_task(1),
                     scheduler.task_manager_for(1).get_task(2),
                     scheduler.task_manager_for(1).get_task(3)};

  auto &thread_state_1 = scheduler.thread_state_for(0);
  auto &task_manager_1 = thread_state_1.get_task_manager();
  auto &thread_state_2 = scheduler.thread_state_for(1);
  auto &task_manager_2 = thread_state_2.get_task_manager();

  SECTION("Must start empty") {
    REQUIRE(!task_manager_1.pop_local_task());
    REQUIRE(!task_manager_1.pop_local_task());
  }

  SECTION("Local push/pop") {
    task_manager_1.push_local_task(tasks_1[0]);
    REQUIRE(task_manager_1.pop_local_task() == tasks_1[0]);
    REQUIRE(!task_manager_1.pop_local_task());
  }

  SECTION("Local push, external pop") {
    task_manager_1.push_local_task(tasks_1[0]);
    REQUIRE(std::get<0>(task_manager_1.steal_task(thread_state_2)) == tasks_1[0]);
    REQUIRE(task_manager_2.pop_clean_task_chain(tasks_1[0]) == tasks_2[0]);
    REQUIRE(task_manager_1.pop_local_task() == nullptr);
  }

  SECTION("Keeps push/pop order #1") {
    task_manager_1.push_local_task(tasks_1[0]);
    task_manager_1.push_local_task(tasks_1[1]);
    REQUIRE(task_manager_1.pop_local_task() == tasks_1[1]);
    REQUIRE(task_manager_1.pop_local_task() == tasks_1[0]);
    REQUIRE(!task_manager_1.pop_local_task());
  }

  SECTION("Keeps push/pop order #2") {
    task_manager_1.push_local_task(tasks_1[0]);
    task_manager_1.push_local_task(tasks_1[1]);
    REQUIRE(std::get<0>(task_manager_1.steal_task(thread_state_2)) == tasks_1[0]);
    REQUIRE(std::get<0>(task_manager_1.steal_task(thread_state_2)) == tasks_1[1]);
  }

  SECTION("Interwined execution #1") {
    // Two top poppers
    task_manager_1.push_local_task(tasks_1[0]);
    REQUIRE(std::get<0>(task_manager_1.steal_task(thread_state_2)) == tasks_1[0]);
    REQUIRE(std::get<0>(task_manager_1.steal_task(thread_state_2)) == nullptr);
  }

  SECTION("Interwined execution #2") {
    // Top and bottom access
    task_manager_1.push_local_task(tasks_1[0]);
    REQUIRE(task_manager_1.pop_local_task() == tasks_1[0]);
    REQUIRE(std::get<0>(task_manager_1.steal_task(thread_state_2)) == nullptr);
  }
}
#endif // PLS_DEQUE_VARIANT == PLS_DEQUE_LOCK_FREE