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

#include "pls/internal/base/error_handling.h"
#include "pls/internal/scheduling/thread_state.h"
#include "pls/internal/scheduling/scheduler.h"

namespace pls::internal::scheduling {

// Get item from locally owned items
strain_local_resource::item_handle strain_local_resource::get_item(unsigned depth) {
  // Only change our resource usage when synced.
  // TODO: maybe relax this requirement
  scheduler::sync();
  thread_state &my_state = thread_state::get();
  auto *active_task = my_state.get_active_task();

  // Find the item to use and acquire it from our local storage
  auto &local_resource = local_items_[my_state.get_thread_id()][depth];
  strain_resource *result = local_resource.resource_;

  // Add it to the task used resources
  result->next_.store(active_task->attached_resources_.load(std::memory_order_relaxed), std::memory_order_relaxed);
  active_task->attached_resources_.store(result, std::memory_order_relaxed);

  // Wrap it for RAII usage on stack
  return strain_local_resource::item_handle{result};
}

strain_local_resource::item_handle::item_handle(strain_resource *resource) : resource_{resource} {
  PLS_ASSERT(resource_->used_.fetch_add(1, std::memory_order_relaxed) == 0,
             "Must not create a handle of a already used resource!");
}

// Return item to locally owned items
strain_local_resource::item_handle::~item_handle() {
  // Only change our resource usage when synced.
  // TODO: maybe relax this requirement
  scheduler::sync();
  thread_state &my_state = thread_state::get();
  auto *active_task = my_state.get_active_task();

  // Remove from task used resources
  strain_resource *previous_resource = nullptr;
  strain_resource *current_resource = active_task->attached_resources_.load(std::memory_order_relaxed);
  while (current_resource != resource_) {
    previous_resource = current_resource;
    current_resource = current_resource->next_.load(std::memory_order_relaxed);
  }

  auto *next_resource = current_resource->next_.load(std::memory_order_relaxed);
  if (previous_resource) {
    previous_resource->next_.store(next_resource, std::memory_order_relaxed);
  } else {
    active_task->attached_resources_.store(next_resource, std::memory_order_relaxed);
  }

  // Give the resource handle back to our local resource array
  auto &local_resource = resource_->strain_local_resource_->local_items_[my_state.get_thread_id()][resource_->depth_];
  local_resource.resource_ = resource_;
  PLS_ASSERT(resource_->used_.fetch_sub(1, std::memory_order_relaxed) == 1,
             "Accidentally freed resource that was accessed multiple times!");
}

strain_resource *strain_local_resource::get_local_copy(strain_resource *other_resources, unsigned thread_id) {
  strain_resource *result = nullptr;

  while (other_resources != nullptr) {
    local_item &local = other_resources->strain_local_resource_->local_items_[thread_id][other_resources->depth_];
    local.resource_->next_.store(result, std::memory_order_relaxed);
    result = local.resource_;

    other_resources = other_resources->next_.load(std::memory_order_relaxed);
  }

  return result;
}

void strain_local_resource::acquire_locally(strain_resource *other_resources, unsigned thread_id) {
  while (other_resources != nullptr) {
    local_item &local = other_resources->strain_local_resource_->local_items_[thread_id][other_resources->depth_];
    local.resource_ = other_resources;

    other_resources = other_resources->next_.load(std::memory_order_relaxed);
  }
}

}