
#ifndef PLS_TASK_MANAGER_H_
#define PLS_TASK_MANAGER_H_

#include <memory>
#include <utility>
#include <array>
#include <mutex>
#include <atomic>

#include "pls/internal/scheduling/task.h"
#include "pls/internal/scheduling/cont_manager.h"
#include "pls/internal/scheduling/memory_block.h"

#include "pls/internal/data_structures/bounded_trading_deque.h"
#include "pls/internal/data_structures/stamped_integer.h"
#include "pls/internal/data_structures/optional.h"

#include "pls/internal/base/spin_lock.h"

namespace pls {
namespace internal {
namespace scheduling {

/**
 * Handles management of tasks in the system. Each thread has a local task manager,
 * responsible for allocating, freeing and publishing tasks for stealing.
 */
class task_manager {
 public:
  // Publishes a task on the stack, i.e. makes it visible for other threads to steal.
  void publish_task(base_task *task) {
    task_deque_.push_bot(task->get_cont()->get_memory_block());
  }

  // Try to pop a local task from this task managers stack.
  data_structures::optional<memory_block *> try_pop_local() {
    return task_deque_.pop_bot().traded_;
  }

  // Try to steal a task from a remote task_manager instance. The stolen task must be stored locally.
  // Returns a pair containing the actual task and if the steal was successful.
  base_task *steal_remote_task(cont_manager &stealing_cont_manager) {
    auto peek = task_deque_.peek_top();

    if (std::get<0>(peek)) {
      memory_block *peeked_memory_block = (*std::get<0>(peek));
      auto peeked_depth = peeked_memory_block->get_depth();

      stealing_cont_manager.move_active_node(peeked_depth);
      auto offered_chain = stealing_cont_manager.get_active_node();
      stealing_cont_manager.move_active_node(1);

      auto stolen_memory_block = task_deque_.pop_top(offered_chain, std::get<1>(peek));
      if (stolen_memory_block) {
        PLS_ASSERT(*stolen_memory_block == peeked_memory_block, "Steal must only work if it is equal!");

        return (*stolen_memory_block)->get_cont()->get_task();
      } else {
        stealing_cont_manager.move_active_node(-(peeked_depth + 1));
        return nullptr;
      }
    }

    return nullptr;
  }

  explicit task_manager(data_structures::bounded_trading_deque<memory_block, memory_block> &task_deque) :
      task_deque_{task_deque} {}

 private:
  data_structures::bounded_trading_deque<memory_block, memory_block> &task_deque_;
};

template<size_t NUM_TASKS>
class static_task_manager {
 public:
  static_task_manager() : task_deque_{}, task_manager_{task_deque_.get_deque()} {};
  task_manager &get_task_manager() { return task_manager_; }

 private:
  data_structures::static_bounded_trading_deque<memory_block, memory_block, NUM_TASKS> task_deque_;
  task_manager task_manager_;
};

}
}
}
#endif //PLS_TASK_MANAGER_H_
