#include <tuple>

#include "pls/internal/scheduling/task_manager.h"

#include "pls/internal/scheduling/task.h"
#include "pls/internal/scheduling/thread_state.h"

namespace pls {
namespace internal {
namespace scheduling {

task_manager::task_manager(task *tasks,
                           data_structures::aligned_stack static_stack_space,
                           size_t num_tasks,
                           size_t stack_size,
                           external_trading_deque &deque) : num_tasks_{num_tasks},
                                                            this_thread_tasks_{tasks},
                                                            active_task_{&tasks[0]},
                                                            deque_{deque} {
  for (size_t i = 0; i < num_tasks - 1; i++) {
    tasks[i].init(static_stack_space.push_bytes(stack_size), stack_size, i, 0);
    if (i > 0) {
      tasks[i].prev_ = &tasks[i - 1];
    }
    if (i < num_tasks - 2) {
      tasks[i].next_ = &tasks[i + 1];
    }
  }
}

static task *find_task(unsigned id, unsigned depth) {
  return thread_state::get().get_scheduler().thread_state_for(id).get_task_manager().get_this_thread_task(depth);
}

void task_manager::push_resource_on_task(task *target_task, task *spare_task_chain) {
  data_structures::stamped_integer current_root;
  data_structures::stamped_integer target_root;
  do {
    current_root = target_task->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_ = 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, target_task->depth_);
      spare_task_chain->resource_stack_next_ = current_root_task;
    }

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

task *task_manager::steal_task(task_manager &stealing_task_manager) {
  PLS_ASSERT(stealing_task_manager.active_task_->depth_ == 0, "Must only steal with clean task chain.");

  auto peek = deque_.peek_top();
  auto optional_target_task = std::get<0>(peek);
  auto target_top = std::get<1>(peek);

  if (optional_target_task) {
    task *target_task = *optional_target_task;
    task *traded_task = stealing_task_manager.active_task_;
    for (unsigned i = 0; i < target_task->depth_; i++) {
      traded_task = traded_task->next_;
    }

    auto optional_result_task = deque_.pop_top(traded_task, target_top);
    if (optional_result_task) {
      return *optional_result_task;
    } else {
      return nullptr;
    }
  } else {
    return nullptr;
  }
}

task *task_manager::pop_resource_from_task(task *target_task) {
  data_structures::stamped_integer current_root;
  data_structures::stamped_integer target_root;
  task *output_task;
  do {
    current_root = target_task->resource_stack_root_.load();
    target_root.stamp = current_root.stamp + 1;

    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, target_task->depth_);
      target_root.value = current_root_task->next_ != nullptr ? current_root_task->next_->thread_id_ + 1 : 0;

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

  return output_task;
}

}
}
}