
#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 continuation = active_task_->next_->run_as_task([lambda, this](context_switcher::continuation cont) {
    auto *last_task = active_task_;
    auto *this_task = active_task_->next_;

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

    traded_cas_field expected_cas_value = deque_.push_bot(active_task_);
    traded_cas_field empty_cas;

    lambda();

    if (active_task_->traded_field_.compare_exchange_strong(expected_cas_value, empty_cas)) {
      // Fast path, simply continue execution where we left of before spawn.
      // This requires no coordination with the resource stack.
      active_task_ = last_task;
      deque_.popped_bot();
      return std::move(last_task->continuation_);
    } else {
      // Slow path, the continuation was stolen.
      // First empty our own deque (everything below must have been stolen already).
      deque_.empty_deque();

      // TODO: This whole process can be taken out into a function, no need to copy it in every lambda.
      //       Also split smaller functions out, this is a mess right now...
      // Try to get a clean resource chain to go back to the main stealing loop
      task *clean_chain = pop_resource_from_task(last_task);
      if (clean_chain == nullptr) {
        // double-check if we are really last one or we only have unlucky timing
        auto cas_field = last_task->traded_field_.load();
        if (cas_field.is_filled_with_object()) {
          traded_cas_field empty_target;
          if (last_task->traded_field_.compare_exchange_strong(cas_field, empty_target)) {
            clean_chain = cas_field.get_trade_object();
          } else {
            clean_chain = pop_resource_from_task(last_task);
          }
        }
      }

      if (clean_chain != nullptr) {
        // We got a clean chain to continue working on.
        PLS_ASSERT(active_task_->prev_->depth_ == clean_chain->depth_,
                   "Resources must only reside in the correct depth!");
        active_task_->prev_ = clean_chain;
        // Walk back chain to make first task active
        while (active_task_->prev_ != nullptr) {
          active_task_ = active_task_->prev_;
        }

        // jump back to continuation in main scheduling loop, time to steal some work
        return std::move(thread_state::get().get_main_continuation());
      } else {
        // We are the last one working on this task. Thus the sync must be finished, continue working.
        active_task_ = last_task;
        return std::move(last_task->continuation_);
      }
    }

    PLS_ERROR("Slow Path/Stealing not implemented!");
  });

  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_
