
#ifndef PLS_TASK_MANAGER_IMPL_H_
#define PLS_TASK_MANAGER_IMPL_H_

#include <memory>
#include <utility>
#include <array>

#include "context_switcher/continuation.h"

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

namespace pls {
namespace internal {
namespace scheduling {

template<typename F>
void task_manager::spawn_child(F &&lambda) {
  auto *spawning_task_manager = this;
  auto continuation =
      active_task_->next_->run_as_task([lambda, spawning_task_manager](context_switcher::continuation cont) {
        auto *last_task = spawning_task_manager->active_task_;
        auto *this_task = spawning_task_manager->active_task_->next_;

        last_task->continuation_ = std::move(cont);
        spawning_task_manager->active_task_ = this_task;

        spawning_task_manager->deque_.push_bot(last_task);
        lambda();
        auto *syncing_task_manager = &thread_state::get().get_task_manager();

        auto pop_result = syncing_task_manager->deque_.pop_bot();
        if (pop_result) {
          // Fast path, simply continue execution where we left of before spawn.
          PLS_ASSERT(*pop_result == last_task,
                     "Fast path, nothing can have changed until here.");
          PLS_ASSERT(spawning_task_manager == syncing_task_manager,
                     "Fast path, nothing can have changed here.");
          PLS_ASSERT(last_task->continuation_.valid(),
                     "Fast path, no one can have continued working on the last task.");

          syncing_task_manager->active_task_ = last_task;
          return std::move(last_task->continuation_);
        } else {
          // Slow path, the continuation was stolen.
          context_switcher::continuation result_cont;
          if (syncing_task_manager->try_clean_return(result_cont)) {
            // We return back to the main scheduling loop
            PLS_ASSERT(result_cont.valid(), "Must only return valid continuations...");
            return std::move(result_cont);
          } else {
            // We finish up the last task and are the sole owner again
            PLS_ASSERT(result_cont.valid(), "Must only return valid continuations...");
            return std::move(result_cont);
          }
        }
      });

  if (continuation.valid()) {
    // We jumped in here from the main loop, keep track!
    thread_state::get().set_main_continuation(std::move(continuation));
  }
}

}
}
}

#endif //PLS_TASK_MANAGER_IMPL_H_
