
#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/scheduling/thread_state.h"

namespace pls {
namespace internal {
namespace scheduling {

/**
 * A task to be executed by the runtime system.
 * Tasks are guaranteed to be executed exactly once.
 *
 * Override the execute_internal() method for your custom code.
 *
 * IMPORTANT:
 * Tasks memory is re-used without calling the destructor.
 * You must call it yourself at the end of execute_internal().
 * This is done to not introduce any overhead of virtual function calls
 * if no clean up is required.
 */
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_offset deque_offset_;

 protected:
  /*
   * Must call the parent constructor.
   *
   * IMPORTANT:
   * Tasks memory is re-used without calling the destructor.
   * You must call it yourself at the end of execute_internal().
   */
  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.
   *
   * Memory is fully self managed. Calling e.g. deconstructors when not needing objects
   * anymore is the users responsibility (memory is simply re-used after the life time of the task ends).
   *
   * @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
  auto item = thread_state::get()->deque_.push_task<T>(std::forward<ARGS>(args)...);

  // Assign forced values (for stack and parent management)
  item->parent_ = this;
  item->finished_construction_ = true;
  item->deque_offset_ = thread_state::get()->deque_.save_offset();

  // Make new task visible to others
  thread_state::get()->deque_.publish_last_task();
}

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!");

  spawn_child<T>(std::forward<ARGS>(args)...);
  // TODO: Check why 'direct spawn' (even when pushing it onto the tas queue) seems to be slower
  //       (Also check if it even is slower or if it only appears so on our laptop)
//  // Push on our deque
//  auto task = thread_state::get()->deque_.push_task<T>(std::forward<ARGS>(args)...);
//
//  // Assign forced values (for stack and parent management)
//  task->parent_ = nullptr; // ...do not assign this to a parent => it will not notify our reference counter
//  task->finished_construction_ = true;
//  task->deque_offset_ = thread_state::get()->deque_.save_offset();
//
//  // Execute it
//  task->execute();

  // Wait for the rest of the tasks
  wait_for_all();
}

}
}
}

#endif //PLS_TASK_H
