
#ifndef PLS_SCHEDULER_IMPL_H
#define PLS_SCHEDULER_IMPL_H

#include <utility>

#include "context_switcher/context_switcher.h"
#include "context_switcher/continuation.h"

#include "pls/internal/scheduling/task_manager.h"
#include "pls/internal/scheduling/task.h"

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

namespace pls {
namespace internal {
namespace scheduling {

class scheduler::init_function {
 public:
  virtual void run() = 0;
};
template<typename F>
class scheduler::init_function_impl : public init_function {
 public:
  explicit init_function_impl(F &function) : function_{function} {}
  void run() override {
    auto &root_task = thread_state::get().get_task_manager().get_active_task();
    root_task.run_as_task([&](context_switcher::continuation cont) {
      thread_state::get().main_continuation() = std::move(cont);
      function_();
      thread_state::get().get_scheduler().work_section_done_.store(true);
      PLS_ASSERT(thread_state::get().main_continuation().valid(), "Must return valid continuation from main task.");
      return std::move(thread_state::get().main_continuation());
    });

  }
 private:
  F &function_;
};

template<typename Function>
void scheduler::perform_work(Function work_section) {
  // Prepare main root task
  init_function_impl<Function> starter_function{work_section};
  main_thread_starter_function_ = &starter_function;

  work_section_done_ = false;
  if (reuse_thread_) {
    auto &my_state = memory_.thread_state_for(0);
    base::this_thread::set_state(&my_state); // Make THIS THREAD become the main worker

    sync_barrier_.wait(); // Trigger threads to wake up
    work_thread_work_section(); // Simply also perform the work section on the main loop
    sync_barrier_.wait(); // Wait for threads to finish
  } else {
    // Simply trigger the others to do the work, this thread will sleep/wait for the time being
    sync_barrier_.wait(); // Trigger threads to wake up
    sync_barrier_.wait(); // Wait for threads to finish
  }
}

template<typename Function>
void scheduler::spawn(Function &&lambda) {
  thread_state::get().get_task_manager().spawn_child(std::forward<Function>(lambda));
}

}
}
}

#endif //PLS_SCHEDULER_IMPL_H
