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

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

task *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) {
  num_resources_++;

  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;
  int iteration = 0;
  do {
    iteration++;
    current_root = this->resource_stack_root_.load();
    target_root.stamp_ = current_root.stamp_ + 1;
    target_root.value_ = spare_task_chain->thread_id_ + 1;

    // TODO: Setting the resource stack next AFTER publishing the task to the CAS field
    //       is a race, as the resource stack next field can be tampered with.
    if (current_root.value_ == 0) {
      // Empty, simply push in with no successor.
      // We are sure that a spare_task_chain is not in any stack when pushing it.
      // Thus, its resource_stack_next_ field must be nullptr.
      // TODO: this should be our race?
      // TODO: add more checks (see if this is violated AND cas succeeds)
      auto *old_value = spare_task_chain->resource_stack_next_.exchange(nullptr);
      if (old_value != nullptr) {
        printf("Why would this invariant happen?\n");
      }
    } 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));
}

void task::reset_task_chain(task *expected_content) {
  num_resources_--;

  data_structures::stamped_integer current_root = this->resource_stack_root_.load();
  PLS_ASSERT(current_root.value_ == expected_content->thread_id_ + 1,
             "Must only reset the task chain if we exactly know its state! (current_root.value_)");

  auto *current_root_task = find_task(current_root.value_ - 1, this->depth_);
  if (current_root_task->resource_stack_next_.load(std::memory_order_relaxed) != nullptr) {
    printf("This could have been the bug...\n");
  }

  data_structures::stamped_integer target_root;
  target_root.stamp_ = current_root.stamp_ + 1;
  bool success = this->resource_stack_root_.compare_exchange_strong(current_root, target_root);
  PLS_ASSERT(success, "Must always succeed in resetting the chain, as we must be the sole one operating on it!");
}

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(num_resources_.fetch_add(-1) > 0, "Must only return an task from the chain if there are items!");

  output_task->resource_stack_next_.store(nullptr);
  return output_task;
}

}
