
#ifndef PLS_TASK_H
#define PLS_TASK_H

#include "pls/internal/helpers/profiler.h"

#include "pls/internal/data_structures/aligned_stack.h"
#include "pls/internal/data_structures/deque.h"

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

namespace pls {
namespace internal {
namespace scheduling {

class task {
  friend class scheduler;

  // Coordinate finishing of sub_tasks
  std::atomic<unsigned int> ref_count_;
  task *parent_;

  // Stack Management (reset stack pointer after wait_for_all() calls)
  data_structures::deque<task>::state deque_state_;

 protected:
  explicit task();

  /**
   * Overwrite this with the actual behaviour of concrete tasks.
   */
  virtual void execute_internal() = 0;

  template<typename T, typename ...ARGS>
  void spawn_child(ARGS &&... args);
  template<typename T, typename ...ARGS>
  void spawn_child_and_wait(ARGS &&... args);
  void wait_for_all();

 private:
  void execute();
};

template<typename T, typename ...ARGS>
void task::spawn_child(ARGS &&... args) {
  PROFILE_FORK_JOIN_STEALING("spawn_child")
  static_assert(std::is_base_of<task, typename std::remove_reference<T>::type>::value, "Only pass task subclasses!");

  // Keep our refcount up to date
  ref_count_++;

  // Push on our deque
  thread_state::get()->deque_.push_tail_cb<T>([this](T *item) {
    // Assign forced values (for stack and parent management)
    item->parent_ = this;
    item->deque_state_ = thread_state::get()->deque_.save_state();
  }, std::forward<ARGS>(args)...);
}

template<typename T, typename ...ARGS>
void task::spawn_child_and_wait(ARGS &&... args) {
  PROFILE_FORK_JOIN_STEALING("spawn_child_wait")
  static_assert(std::is_base_of<task, typename std::remove_reference<T>::type>::value, "Only pass task subclasses!");

  // Assign forced values (for stack and parent management)
  // TODO: Move this after construction
  T sub_task{std::forward<ARGS>(args)...};
  sub_task.parent_ = nullptr;
  sub_task.deque_state_ = thread_state::get()->deque_.save_state();
  PROFILE_END_BLOCK
  sub_task.execute();

  wait_for_all();
}

}
}
}

#endif //PLS_TASK_H
