
#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;

  // Memory-Management (allow to allocate memory blocks in constructor)
  bool finished_construction_;

  // 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:
  /*
   * Must call the parent constructor.
   */
  explicit task();

  /**
   * Allow to allocate extra memory during run-time for this task.
   * Memory will be pushed onto the stack (in aligned memory, thus avoid many small chunks).
   * MUST be called in constructor, never afterwards.
   *
   * @param size Number of bytes to be allocated
   * @return The allocated memory region
   */
  void *allocate_memory(long size);

  /**
   * 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->finished_construction_ = true;
    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) {
  static_assert(std::is_base_of<task, typename std::remove_reference<T>::type>::value, "Only pass task subclasses!");

  // TODO: See if we can inline this (avoid counters/deque) while maintaining memory management
  spawn_child<T>(std::forward<ARGS>(args)...);
  wait_for_all();
}

}
}
}

#endif //PLS_TASK_H
