#include <pls/internal/base/backoff.h>
#include "pls/internal/helpers/profiler.h"

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

namespace pls {
namespace internal {
namespace scheduling {

bool abstract_task::steal_work() {
  thread_local static base::backoff backoff{};

  PROFILE_STEALING("abstract_task::steal_work")
  const auto my_state = base::this_thread::state<thread_state>();
  const auto my_scheduler = my_state->scheduler_;

  const size_t my_id = my_state->id_;
  const size_t offset = my_state->random_() % my_scheduler->num_threads();
  const size_t max_tries = my_scheduler->num_threads() - 1; // TODO: Tune this value
  for (size_t i = 0; i < max_tries; i++) {
    size_t target = (offset + i) % my_scheduler->num_threads();
    if (target == my_id) {
      target = (target + 1) % my_scheduler->num_threads();
    }
    auto target_state = my_scheduler->thread_state_for(target);

    if (!target_state->lock_.reader_try_lock()) {
      continue;
    }

    // Dig down to our level
    PROFILE_STEALING("Go to our level")
    abstract_task *current_task = target_state->root_task_;
    while (current_task != nullptr && current_task->depth() < depth()) {
      current_task = current_task->child();
    }
    PROFILE_END_BLOCK

    // Try to steal 'internal', e.g. for_join_sub_tasks in a fork_join_task constellation
    PROFILE_STEALING("Internal Steal")
    if (current_task != nullptr) {
      // See if it equals our type and depth of task
      if (current_task->unique_id_ == unique_id_ &&
          current_task->depth_ == depth_) {
        if (internal_stealing(current_task)) {
          // internal steal was a success, hand it back to the internal scheduler
          target_state->lock_.reader_unlock();
          backoff.reset();
          return true;
        }

        // No success, we need to steal work from a deeper level using 'top level task stealing'
        current_task = current_task->child();
      }
    }
    PROFILE_END_BLOCK;


    // Execute 'top level task steal' if possible
    // (only try deeper tasks to keep depth restricted stealing).
    PROFILE_STEALING("Top Level Steal")
    while (current_task != nullptr) {
      auto lock = &target_state->lock_;
      if (current_task->split_task(lock)) {
        // top level steal was a success (we did a top level task steal)
        backoff.reset();
        return false;
      }

      current_task = current_task->child_task_;
    }
    PROFILE_END_BLOCK;
    target_state->lock_.reader_unlock();
  }

  // internal steal was no success
  backoff.do_backoff();
  return false;
}

}
}
}