
#ifndef PLS_WORK_STEALING_DEQUE_H_
#define PLS_WORK_STEALING_DEQUE_H_

#include <atomic>
#include <mutex>
#include <pls/internal/scheduling/thread_state.h>

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

#include "aligned_stack.h"

namespace pls {
namespace internal {
namespace data_structures {

using cas_integer = base::system_details::cas_integer;
using pointer_t = base::system_details::pointer_t;
static cas_integer get_stamp(cas_integer n) {
  return (n & base::system_details::FIRST_HALF_CAS_INTEGER) >> ((sizeof(cas_integer) / 2) * 8);
}
static cas_integer get_offset(cas_integer n) {
  return n & base::system_details::SECOND_HALF_CAS_INTEGER;
}
static cas_integer set_stamp(cas_integer n, cas_integer new_value) {
  return (new_value << ((sizeof(cas_integer) / 2) * 8)) | (n & base::system_details::SECOND_HALF_CAS_INTEGER);
}
//static cas_integer set_offset(cas_integer n, cas_integer new_value) {
//  return new_value | (n & base::system_details::FIRST_HALF_CAS_INTEGER);
//}

class work_stealing_deque_item {
  // Pointer to the actual data
  pointer_t data_;
  // Index (relative to stack base) to the next and previous element
  cas_integer next_item_;
  cas_integer previous_item_;

 public:
  work_stealing_deque_item() : data_{0}, next_item_{0}, previous_item_{0} {}

  template<typename Item>
  Item *data() {
    return reinterpret_cast<Item *>(data_);
  }

  template<typename Item>
  void set_data(Item *data) {
    data_ = reinterpret_cast<pointer_t >(data);
  }

  cas_integer next_item() {
    return next_item_;
  }
  void set_next_item(cas_integer next_item) {
    next_item_ = next_item;
  }
  cas_integer previous_item() {
    return previous_item_;
  }
  void set_previous_item(cas_integer previous_item) {
    previous_item_ = previous_item;
  }
};
static_assert(sizeof(work_stealing_deque_item) < base::system_details::CACHE_LINE_SIZE,
              "Work stealing deque relies on memory layout and requires cache lines to be longer than one 'work_stealing_deque_item' instance!");

template<typename Item>
class work_stealing_deque {
  // Deque 'takes over' stack and handles memory management while in use.
  // At any point in time the deque can stop using more memory and the stack can be used by other entities.
  aligned_stack *stack_;
  pointer_t base_pointer_;

  std::atomic<cas_integer> head_;
  std::atomic<cas_integer> tail_;
  cas_integer previous_tail_;

  base::spin_lock lock_{}; // TODO: Remove after debugging


 public:
  using state = aligned_stack::state;

  explicit work_stealing_deque(aligned_stack *stack) : stack_{stack},
                                                       base_pointer_{0},
                                                       head_{0},
                                                       tail_{0},
                                                       previous_tail_{0} {
    reset_base_pointer();
  }
  work_stealing_deque(const work_stealing_deque &other) : stack_{other.stack_},
                                                          base_pointer_{other.base_pointer_},
                                                          head_{other.head_.load()},
                                                          tail_{other.tail_.load()},
                                                          previous_tail_{other.previous_tail_} {}

  void reset_base_pointer() {
    base_pointer_ = reinterpret_cast<pointer_t >(stack_->save_state()); // Keep the base of our region in the stack
  }

  work_stealing_deque_item *item_at(cas_integer position) {
    return reinterpret_cast<work_stealing_deque_item *>(base_pointer_
        + (base::system_details::CACHE_LINE_SIZE * position));
  }

  cas_integer current_stack_offset() {
    return (stack_->save_state() - base_pointer_) / base::system_details::CACHE_LINE_SIZE;
  }

  template<typename T>
  std::pair<work_stealing_deque_item, T> *allocate_item(const T &new_item) {
    // 'Union' type to push both on stack
    using pair_t = std::pair<work_stealing_deque_item, T>;
    // Allocate space on stack
    auto new_pair = reinterpret_cast<pair_t *>(stack_->push<pair_t>());
    // Initialize memory on stack
    new((void *) &(new_pair->first)) work_stealing_deque_item();
    new((void *) &(new_pair->second)) T(new_item);

    return new_pair;
  }

  template<typename T>
  Item *push_tail(const T &new_item) {
    cas_integer local_tail = tail_;

    auto new_pair = allocate_item(new_item);
    // Prepare current tail to point to correct next items
    auto tail_deque_item = item_at(local_tail);
    tail_deque_item->set_data(&(new_pair->second));
    tail_deque_item->set_next_item(current_stack_offset());
    tail_deque_item->set_previous_item(previous_tail_);
    previous_tail_ = local_tail;

    // Linearization point, item appears after this write
    cas_integer new_tail = current_stack_offset();
    tail_ = new_tail;

    return &(new_pair->second);
  }

  Item *pop_tail() {
    cas_integer local_tail = tail_;
    cas_integer local_head = head_;

    if (local_tail <= get_offset(local_head)) {
      return nullptr; // EMPTY
    }

    work_stealing_deque_item *previous_tail_item = item_at(previous_tail_);
    cas_integer new_tail = previous_tail_;
    previous_tail_ = previous_tail_item->previous_item();

    // Publish our wish to set the tail back
    tail_ = new_tail;
    // Get the state of local head AFTER we published our wish
    local_head = head_; // Linearization point, outside knows list is empty

    if (get_offset(local_head) < new_tail) {
      return previous_tail_item->data<Item>(); // Success, enough distance to other threads
    }

    if (get_offset(local_head) == new_tail) {
      cas_integer new_head = set_stamp(new_tail, get_stamp(local_head) + 1);
      // Try competing with consumers by updating the head's stamp value
      if (head_.compare_exchange_strong(local_head, new_head)) {
        return previous_tail_item->data<Item>(); // SUCCESS, we won the competition with other threads
      }
    }

    // Some other thread either won the competition or it already set the head further than we are
    // before we even tried to compete with it.
    // Reset the queue into an empty state => head_ = tail_
    tail_ = get_offset(local_head); // ...we give up to the other winning thread

    return nullptr; // EMPTY, we lost the competition with other threads
  }

  Item *pop_head() {
    cas_integer local_head = head_;
    cas_integer local_tail = tail_;

    if (local_tail <= get_offset(local_head)) {
      return nullptr; // EMPTY
    }
    // Load info on current deque item.
    // In case we have a race with a new (aba) overwritten item at this position,
    // there has to be a competition over the tail -> the stamp increased and our next
    // operation will fail anyways!
    work_stealing_deque_item *head_deque_item = item_at(get_offset(local_head));
    cas_integer next_item_offset = head_deque_item->next_item();
    Item *head_data_item = head_deque_item->data<Item>();

    // We try to set the head to this new position.
    // Possible outcomes:
    // 1) no one interrupted us, we win this competition
    // 2) other thread took the head, we lose to this
    // 3) owning thread removed tail, we lose to this
    cas_integer new_head = set_stamp(next_item_offset, get_stamp(local_head) + 1);
    if (head_.compare_exchange_strong(local_head, new_head)) {
      return head_data_item; // SUCCESS, we won the competition
    }

    return nullptr; // EMPTY, we lost the competition

  }

  void release_memory_until(state state) {
    cas_integer item_offset = (state - base_pointer_) / base::system_details::CACHE_LINE_SIZE;

    cas_integer local_head = head_;
    cas_integer local_tail = tail_;

    stack_->reset_state(state);

    if (item_offset < local_tail) {
      tail_ = item_offset;
      if (get_offset(local_head) >= local_tail) {
        head_ = set_stamp(item_offset, get_stamp(local_head) + 1);
      }
    }
  }

  void release_memory_until(Item *item) {
    release_memory_until(reinterpret_cast<pointer_t >(item));
  }

  state save_state() {
    return stack_->save_state();
  }
};

}
}
}

#endif //PLS_WORK_STEALING_DEQUE_H_
