
#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 *last_task = spawning_task_manager->active_task_;
  auto *spawned_task = spawning_task_manager->active_task_->next_;

  auto continuation =
      spawned_task->run_as_task([=](context_switcher::continuation cont) {
        // allow stealing threads to continue the last task.
        last_task->continuation_ = std::move(cont);

        // we are now executing the new task, allow others to steal the last task continuation.
        spawned_task->is_synchronized_ = true;
        spawning_task_manager->active_task_ = spawned_task;
        spawning_task_manager->deque_.push_bot(last_task);

        // execute the lambda itself, which could lead to a different thread returning.
        lambda();
        auto *syncing_task_manager = &thread_state::get().get_task_manager();
        PLS_ASSERT(syncing_task_manager->active_task_ == spawned_task,
                   "Task manager must always point its active task onto whats executing.");

        // try to pop a task of the syncing task manager.
        // possible outcomes:
        // - this is a different task manager, it must have an empty deque and fail
        // - this is the same task manager and someone stole last tasks, thus this will fail
        // - this is the same task manager and no one stole the last task, this this will succeed
        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 last task was stolen. Sync using the resource stack.
          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().main_continuation() = std::move(continuation);
  }
}

}
}
}

#endif //PLS_TASK_MANAGER_IMPL_H_
