
#ifndef PLS_TBB_LIKE_TASK_H
#define PLS_TBB_LIKE_TASK_H

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

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

#include "abstract_task.h"
#include "thread_state.h"

namespace pls {
namespace internal {
namespace scheduling {

class fork_join_task;
class fork_join_sub_task {
  friend class fork_join_task;

  // Coordinate finishing of sub_tasks
  std::atomic_uint32_t ref_count_;
  fork_join_sub_task *parent_;

  // Access to TBB scheduling environment
  fork_join_task *tbb_task_;

  bool executed = false;
  int executed_at = -1;

  // Stack Management (reset stack pointer after wait_for_all() calls)
  data_structures::work_stealing_deque<fork_join_sub_task>::state deque_state_;
 protected:
  explicit fork_join_sub_task();
  fork_join_sub_task(const fork_join_sub_task &other);

  // Overwritten with behaviour of child tasks
  virtual void execute_internal() = 0;

 public:
  // Only use them when actually executing this sub_task (only public for simpler API design)
  template<typename T>
  void spawn_child(T &sub_task);
  void wait_for_all();

  static fork_join_sub_task *current();
 private:
  void execute();
};

template<typename Function>
class fork_join_lambda_by_reference : public fork_join_sub_task {
  const Function &function_;

 public:
  explicit fork_join_lambda_by_reference(const Function &function) : fork_join_sub_task{}, function_{function} {};

 protected:
  void execute_internal() override {
    function_();
  }
};

template<typename Function>
class fork_join_lambda_by_value : public fork_join_sub_task {
  const Function function_;

 public:
  explicit fork_join_lambda_by_value(const Function &function) : fork_join_sub_task{}, function_{function} {};

 protected:
  void execute_internal() override {
    function_();
  }
};

class fork_join_task : public abstract_task {
  friend class fork_join_sub_task;

  fork_join_sub_task *root_task_;
  fork_join_sub_task *currently_executing_;

  // Double-Ended Queue management
  data_structures::work_stealing_deque<fork_join_sub_task> deque_;

  // Steal Management
  fork_join_sub_task *last_stolen_;

  fork_join_sub_task *get_local_sub_task();
  fork_join_sub_task *get_stolen_sub_task();

  bool internal_stealing(abstract_task *other_task) override;
  bool split_task(base::swmr_spin_lock * /*lock*/) override;

 public:
  explicit fork_join_task(fork_join_sub_task *root_task, const abstract_task::id &id);
  void execute() override;
  fork_join_sub_task *currently_executing() const;
};

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

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

  // Assign forced values
  task.parent_ = this;
  task.tbb_task_ = tbb_task_;
  task.deque_state_ = tbb_task_->deque_.save_state();

  // Push on our deque
  const T const_task = task;
  tbb_task_->deque_.push_tail(const_task);
}

}
}
}

#endif //PLS_TBB_LIKE_TASK_H
