#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::propose_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_split_integer current_root;
  data_structures::stamped_split_integer target_root;
  do {
    current_root = this->resource_stack_root_.load();
    PLS_ASSERT(current_root.value_2_ == 0, "Must only propose one push at a time!");

    // Add it to the current stack state by chaining its next stack task to the current base.
    // Popping threads will see this proposed task as if it is the first item in the stack and pop it first,
    // making sure that the chain is valid without any further modification when accepting the proposed task.
    auto *current_root_task = current_root.value_1_ == 0 ? nullptr : find_task(current_root.value_1_ - 1, this->depth_);
    spare_task_chain->resource_stack_next_.store(current_root_task);

    target_root.stamp_ = current_root.stamp_ + 1;
    target_root.value_1_ = current_root.value_1_;
    target_root.value_2_ = spare_task_chain->thread_id_ + 1;
  } while (!this->resource_stack_root_.compare_exchange_strong(current_root, target_root));
}

bool task::accept_proposed() {
  data_structures::stamped_split_integer current_root;
  data_structures::stamped_split_integer target_root;
  do {
    current_root = this->resource_stack_root_.load();
    if (current_root.value_2_ == 0) {
      return false; // We are done, nothing to accept!
    }

    target_root.stamp_ = current_root.stamp_ + 1;
    target_root.value_1_ = current_root.value_2_;
    target_root.value_2_ = 0;
  } while (!this->resource_stack_root_.compare_exchange_strong(current_root, target_root));

  return true;
}

bool task::decline_proposed() {
  bool proposed_still_there;

  data_structures::stamped_split_integer current_root;
  data_structures::stamped_split_integer target_root;
  do {
    current_root = this->resource_stack_root_.load();
    proposed_still_there = current_root.value_2_ != 0;

    // No need to fetch anything, just delete the proposed item.
    target_root.stamp_ = current_root.stamp_ + 1;
    target_root.value_1_ = current_root.value_1_;
    target_root.value_2_ = 0;
  } while (!this->resource_stack_root_.compare_exchange_strong(current_root, target_root));

  return proposed_still_there;
}

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.");
  propose_push_task_chain(spare_task_chain);
  accept_proposed();
}

task *task::pop_task_chain() {
  data_structures::stamped_split_integer current_root;
  data_structures::stamped_split_integer target_root;

  task *output_task;
  do {
    current_root = this->resource_stack_root_.load();

    if (current_root.value_2_ == 0) {
      if (current_root.value_1_ == 0) {
        // No main chain and no proposed element.
        return nullptr;
      } else {
        // No entries in the proposed slot, but fond
        // elements in primary stack, try to pop from there.
        auto *current_root_task = find_task(current_root.value_1_ - 1, this->depth_);
        auto *next_stack_task = current_root_task->resource_stack_next_.load();

        target_root.stamp_ = current_root.stamp_ + 1;
        target_root.value_1_ = next_stack_task != nullptr ? next_stack_task->thread_id_ + 1 : 0;
        target_root.value_2_ = 0;

        output_task = current_root_task;
      }
    } else {
      // We got a proposed element. Treat it as beginning of resource stack by
      // popping it instead of the main element chain.
      auto *proposed_task = find_task(current_root.value_2_ - 1, this->depth_);

      // Empty out the proposed slot
      target_root.stamp_ = current_root.stamp_ + 1;
      target_root.value_1_ = current_root.value_1_;
      target_root.value_2_ = 0;

      output_task = proposed_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;
}

}