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

namespace pls::internal::scheduling::lock_free {

// TODO: this 'global' lookup hardly bound to the full scheduler to be setup could be reworked for better testing.
static task *find_task(unsigned id, unsigned depth) {
  return thread_state::get().get_scheduler().thread_state_for(id).get_task_manager().get_task(depth);
}

void task::push_task_chain(task *spare_task_chain) {
  PLS_ASSERT(this->thread_id_ != spare_task_chain->thread_id_,
             "Makes no sense to push task onto itself, as it is not clean by definition.");
  PLS_ASSERT(this->depth_ == spare_task_chain->depth_,
             "Must only push tasks with correct depth.");

  data_structures::stamped_integer current_root;
  data_structures::stamped_integer target_root;
  do {
    current_root = this->resource_stack_root_.load();
    target_root.stamp = current_root.stamp + 1;
    target_root.value = spare_task_chain->thread_id_ + 1;

    if (current_root.value == 0) {
      // Empty, simply push in with no successor
      spare_task_chain->resource_stack_next_.store(nullptr);
    } else {
      // Already an entry. Find it's corresponding task and set it as our successor.
      auto *current_root_task = find_task(current_root.value - 1, this->depth_);
      spare_task_chain->resource_stack_next_.store(current_root_task);
    }

  } while (!this->resource_stack_root_.compare_exchange_strong(current_root, target_root));
}

task *task::pop_task_chain() {
  data_structures::stamped_integer current_root;
  data_structures::stamped_integer target_root;
  task *output_task;
  do {
    current_root = this->resource_stack_root_.load();
    if (current_root.value == 0) {
      // Empty...
      return nullptr;
    } else {
      // Found something, try to pop it
      auto *current_root_task = find_task(current_root.value - 1, this->depth_);
      auto *next_stack_task = current_root_task->resource_stack_next_.load();

      target_root.stamp = current_root.stamp + 1;
      target_root.value = next_stack_task != nullptr ? next_stack_task->thread_id_ + 1 : 0;

      output_task = current_root_task;
    }
  } while (!this->resource_stack_root_.compare_exchange_strong(current_root, target_root));

  PLS_ASSERT(scheduler::check_task_chain_backward(*output_task), "Must only pop proper task chains.");
  output_task->resource_stack_next_.store(nullptr);
  return output_task;
}

void task::reset_task_chain() {
  auto current_root = this->resource_stack_root_.load();
  current_root.stamp++;
  current_root.value = 0;
  this->resource_stack_root_.store(current_root);
  this->resource_stack_next_.store(nullptr);
}

}