
#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_ws_deque.h"
#include "pls/internal/data_structures/stamped_integer.h"

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

namespace pls {
namespace internal {
namespace scheduling {

struct task_handle {
 public:
  task_handle() : task_{nullptr}, task_memory_block_{nullptr} {};
  explicit task_handle(base_task *task) : task_{task},
                                          task_memory_block_{task->get_cont()->get_memory_block()} {};

  base_task *task_;
  // This seems redundant first, but is needed for a race-free steal.
  // It could happen that the task's memory is overwritten and the pointer to it's memory block gets invalid.
  // We can do this more elegantly in the future.
  memory_block *task_memory_block_;
};

/**
 * 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) {
    std::lock_guard<base::spin_lock> lock{lock_};
    task_deque_.push_bottom(task_handle{task});
  }

  // Try to pop a local task from this task managers stack.
  bool steal_local_task() {
    std::lock_guard<base::spin_lock> lock{lock_};
    return task_deque_.pop_bottom();
  }

  // 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) {
    std::lock_guard<base::spin_lock> lock{lock_};

    // TODO: See if we can somehow make this trade lock free (and still be correct)

    auto stolen_task_handle = task_deque_.pop_top();
    if (stolen_task_handle) {
      base_task *stolen_task = (*stolen_task_handle).task_;
      memory_block *stolen_task_memory = (*stolen_task_handle).task_memory_block_;
      auto stolen_task_depth = stolen_task_memory->get_depth();
      auto &atomic_state = stolen_task_memory->get_state();
      auto &atomic_offered_chain = stolen_task_memory->get_offered_chain();

      // TODO: We ignore all we tried with lock free implementations here, just store the state how it is supposed to be
      stealing_cont_manager.move_active_node(stolen_task_depth);
      auto offered_chain = stealing_cont_manager.get_active_node();
      stealing_cont_manager.move_active_node(1);
      atomic_offered_chain.store(offered_chain);
      atomic_state.store(memory_block::stolen);

      return stolen_task;
    } else {
      return nullptr;
    }
  }

  explicit task_manager(data_structures::bounded_ws_deque<task_handle> &task_deque) : task_deque_{task_deque},
                                                                                      lock_{} {}

 private:
  data_structures::bounded_ws_deque<task_handle> &task_deque_;
  base::spin_lock lock_;
};

template<size_t NUM_TASKS, size_t MAX_STACK_SIZE>
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_ws_deque<task_handle, NUM_TASKS> task_deque_;
  task_manager task_manager_;
};

}
}
}
#endif //PLS_TASK_MANAGER_H_
