
#ifndef PLS_INTERNAL_SCHEDULING_STACK_RESOURCE_H_
#define PLS_INTERNAL_SCHEDULING_STACK_RESOURCE_H_

#include "pls/internal/base/system_details.h"

#include <vector>
#include <atomic>
#include <memory>

namespace pls::internal::scheduling {

class strain_local_resource;
struct PLS_CACHE_ALIGN strain_resource {
  strain_resource(strain_local_resource *resource, unsigned index, unsigned depth) :
      strain_local_resource_{resource}, index_{index}, depth_{depth} {}

  strain_resource(const strain_resource &) = delete;
  strain_resource &operator=(const strain_resource &) = delete;
  strain_resource(strain_resource &&) = delete;
  strain_resource &operator=(strain_resource &&) = delete;

  strain_local_resource *const strain_local_resource_;
  unsigned const index_;
  unsigned const depth_;

  std::atomic<int> used_{0};
  std::atomic<strain_resource *> next_{nullptr};
};

class strain_local_resource {
  struct local_item {
    local_item(strain_local_resource *strain_local_resource, unsigned index, unsigned depth) :
        initial_resource_{std::make_unique<strain_resource>(strain_local_resource, index, depth)},
        resource_{initial_resource_.get()} {}

    std::unique_ptr<strain_resource> initial_resource_;
    strain_resource *resource_;
  };
 public:
  // Handle to a stack-like resource (index of it).
  // MUST be handled in a RAII like manner on the stack
  // and must not be copied or moved for this to work.
  class item_handle {
   public:
    item_handle(const item_handle &) = delete;
    item_handle &operator=(const item_handle &) = delete;
    item_handle(item_handle &&) = delete;
    item_handle &operator=(item_handle &&) = delete;

    explicit item_handle(strain_resource *resource);
    ~item_handle();

    unsigned get_strain_index() {
      return resource_->index_;
    }

   private:
    strain_resource *const resource_;
  };

  strain_local_resource(unsigned num_threads,
                        unsigned depth) : num_threads_{num_threads}, local_items_(num_threads) {
    for (unsigned thread_id = 0; thread_id < num_threads; thread_id++) {
      local_items_[thread_id].reserve(depth);
      for (unsigned i = 0; i < depth; i++) {
        // Start by each thread owning its local items
        local_items_[thread_id].emplace_back(this, thread_id, i);
      }
    }
  }

  [[nodiscard]] unsigned get_num_threads() const { return num_threads_; }
  item_handle get_item(unsigned depth);
  static strain_resource *get_local_copy(strain_resource *other_resources, unsigned thread_id);
  static void acquire_locally(strain_resource *other_resources, unsigned thread_id);

 private:
  const unsigned num_threads_;
  std::vector<std::vector<local_item>> local_items_;
};

}

#endif //PLS_INTERNAL_SCHEDULING_STACK_RESOURCE_H_
