#pragma once
#include "sortlib.cuh"
#ifndef LOCKFREE_SKIPLIST_CUH_
#define LOCKFREE_SKIPLIST_CUH_

#ifndef LOCKFREE_KEYTYPE_
#define LOCKFREE_KEYTYPE_
typedef unsigned long long key_type;
#endif

#ifndef BUILD_SIZE
#define BUILD_SIZE 1048576
#endif

#ifdef MEASURE_TIME
#undef BUILD_SIZE
#define BUILD_SIZE 1024
#endif

#ifndef LOCKFREE_DEFINITIONS_
#define LOCKFREE_DEFINITIONS_
constexpr size_t MAX_LEVEL = 16;

constexpr size_t NUM_ITEMS = BUILD_SIZE;
#endif

class Node;

__device__ Node **nodes;

// Definition of generic node class

class __attribute__((aligned(16))) Node {
public:
  int topLevel;                 // Level of the node
  key_type key;                 // Key value
  key_type next[MAX_LEVEL + 1]; // Array of next links

  // Create a next field from a reference and mark bit
  __device__ __host__ key_type CreateRef(Node *ref, bool mark) {
    auto val = (key_type)ref;
    val = val | mark;
    return val;
  }

  __device__ __host__ void SetRef(int index, Node *ref, bool mark) {
    next[index] = CreateRef(ref, mark);
  }

  // Extract the reference from a next field
  __device__ Node *GetReference(int index) {
    key_type ref = next[index];
    return (Node *)((ref >> 1) << 1);
  }

  // Extract the reference and mark bit from a next field
  __device__ Node *Get(int index, bool *marked) {
    marked[0] = next[index] % 2;
    return (Node *)((next[index] >> 1) << 1);
  }

  // CompareAndSet wrapper
  __device__ bool CompareAndSet(int index, Node *expectedRef, Node *newRef,
                                bool oldMark, bool newMark) {
    key_type oldVal = (key_type)expectedRef | oldMark;
    key_type newVal = (key_type)newRef | newMark;
    key_type *ref = &(next[index]);
    key_type oldValOut = atomicCAS(ref, oldVal, newVal);
    if (oldValOut == oldVal)
      return true;
    return false;
  }

  // Constructor for sentinel nodes
  Node(key_type k) {
    key = k;
    topLevel = MAX_LEVEL;
    int i;
    for (i = 0; i < MAX_LEVEL + 1; i++) {
      next[i] = CreateRef((Node *)nullptr, false);
    }
  }
};

// Definition of lock-free skip list

class LockFreeSkipList {
  static constexpr size_t SAMPLE_LENGTH = 1024;

public:
  Node *head;
  Node *tail;

  unsigned int *pointerIndex;

  key_type *samples;

  CustomSort searcher;

  double slice_size;

  LockFreeSkipList(key_type *samples, size_t sample_length)
      : searcher(sample_length, sizeof(key_type)) {
    Node *h = new Node(0);
    // size_ = 0;
#if __WORDSIZE == 64
    Node *t = new Node((key_type)NUM_ITEMS + 10);
#else
    Node *t = new Node(0xffffffffULL);
#endif
    cudaMalloc(&head, sizeof(Node));

    cudaMalloc(&tail, sizeof(Node));
    int i;
    for (i = 0; i < h->topLevel + 1; i++) {
      h->SetRef(i, tail, false);
    }
    cudaMemcpy(head, h, sizeof(Node), cudaMemcpyHostToDevice);

    cudaMemcpy(tail, t, sizeof(Node), cudaMemcpyHostToDevice);

    cudaMalloc(&pointerIndex, sizeof(unsigned int));
    initDeviceVariable();

    this->slice_size = 1.0 / sample_length;

    cudaMalloc(&samples, sizeof(key_type) * sample_length);
    cudaMemcpy(this->samples, samples, sizeof(key_type) * sample_length,
               cudaMemcpyHostToDevice);
  }

  void initDeviceVariable() const {
    unsigned int x = 0;
    cudaMemcpy(pointerIndex, &x, sizeof(unsigned int), cudaMemcpyHostToDevice);
  }

  ~LockFreeSkipList() { cudaFree(samples); }

  __device__ unsigned static trailing_zeroes(size_t index) {
    constexpr auto block_size = 2;
    unsigned bits = 0;
    unsigned x = index / block_size;

    if (x) {
      while (x % block_size == 0) {
        ++bits;
        x /= block_size;
      }
    }
    return bits;
  }

  __device__ size_t calcLevel(key_type k) {
    auto index =
        this->searcher.sample_cdf_custom_version(this->samples, k) / slice_size;
    auto level = trailing_zeroes(index);
    return level;
  }

  __device__ bool find(key_type, Node **, Node **) const; // Helping method
  __device__ bool Add(key_type);
  //__device__ bool Delete(key_type);
  __device__ bool Search(key_type) const;

#ifdef MEASURE_ACCESS
  unsigned access_times = 0;

  __device__ unsigned getAccessCount() const { return this->access_times; }
  __device__ void increaseAccessCount(unsigned count = 1) {
    atomicAdd(&this->access_times, count);
  }
#else
  __device__ void increaseAccessCount(unsigned _count = 1) {}
#endif

#ifdef MEASURE_TIME
  unsigned round = 0;
  __device__ void increaseRoundCount(unsigned count = 1) {
    atomicAdd(&this->round, count);
  }
  int spend_time[NUM_ITEMS]{0};
  unsigned long long total_time = 0;
  __device__ unsigned getRoundCount() const { return this->round; }
#endif
};

//__device__ Node **nodes;                  // Pool of pre-allocated nodes
//__device__ unsigned int pointerIndex = 0; // Index into pool of free nodes
//__device__ key_type *randoms; // Array storing the levels of the nodes in the
// free pool

// Function for creating a new node when requested by an add operation

__device__ Node *GetNewNode(key_type key, unsigned int *pointerIndex,
                            int topLevel) {
  key_type ind = atomicInc(pointerIndex, NUM_ITEMS);
  Node *n = nodes[ind];
  n->key = key;
  n->topLevel = topLevel;
  for (int i = 0; i < n->topLevel + 1; i++) {
    n->SetRef(i, nullptr, false);
  }
  return n;
}

// Find the window holding key
// On the way clean up logically deleted nodes (those with set marked bit)

__device__ bool LockFreeSkipList::find(
    key_type key, Node **preds,
    Node **succs) const { // preds and succs are arrays of pointers
  int bottomLevel = 0;
  bool marked[] = {false};
  bool snip;
  Node *pred;
  Node *curr;
  Node *succ;
  bool beenThereDoneThat;
  while (true) {
    beenThereDoneThat = false;
    pred = head;
    int level;
    for (level = MAX_LEVEL; level >= bottomLevel; level--) {
      curr = pred->GetReference(level);
      while (true) {
        succ = curr->Get(level, marked);
        while (marked[0]) {
          snip = pred->CompareAndSet(level, curr, succ, false, false);
          beenThereDoneThat = true;
          if (!snip)
            break;
          curr = pred->GetReference(level);
          succ = curr->Get(level, marked);
          beenThereDoneThat = false;
          // printf("find key is %d \n",(int)key);
        }
        if (beenThereDoneThat)
          break;
        if (curr->key <= key) {
          pred = curr;
          curr = succ;
        } else {
          break;
        }
      }
      if (beenThereDoneThat)
        break;
      preds[level] = pred;
      succs[level] = curr;
    }
    if (beenThereDoneThat)
      continue;
    return ((curr->key == key));
  }
}

__device__ bool LockFreeSkipList::Search(key_type key) const {
  int bottomLevel = 0;
  bool marked = false;
  Node *pred = head;
  Node *curr = nullptr;
  Node *succ;
  int level;
  for (level = MAX_LEVEL; level >= bottomLevel; level--) {
    curr = pred->GetReference(level);
#ifdef MEASURE_ACCESS
    this->increaseAccessCount();
#endif
    while (true) {
      succ = curr->Get(level, &marked);
#ifdef MEASURE_ACCESS
      this->increaseAccessCount();
#endif
      while (marked) {
        curr = curr->GetReference(level);
        succ = curr->Get(level, &marked);
#ifdef MEASURE_ACCESS
        this->increaseAccessCount(2);
#endif
      }
      if (curr->key < key) {
        pred = curr;
        curr = succ;
      } else {
        break;
      }
    }
  }
  return (curr != nullptr && curr->key == key);
}

/*__device__ bool LockFreeSkipList::Delete(key_type key) {
  int bottomLevel = 0;
  Node *preds[MAX_LEVEL + 1];
  Node *succs[MAX_LEVEL + 1];
  Node *succ;
  bool marked[] = {false};
  while (true) {
    bool found = find(key, preds, succs);
    if (!found) {
      return false;
    } else {
      Node *nodeToDelete = succs[bottomLevel];
      int level;
      for (level = nodeToDelete->topLevel; level >= bottomLevel + 1; level--) {
        succ = nodeToDelete->Get(level, marked);
        while (!marked[0]) {
          nodeToDelete->CompareAndSet(level, succ, succ, false, true);
          succ = nodeToDelete->Get(level, marked);
        }
      }
      succ = nodeToDelete->Get(bottomLevel, marked);
      while (true) {
        bool iMarkedIt =
            nodeToDelete->CompareAndSet(bottomLevel, succ, succ, false, true);
        succ = succs[bottomLevel]->Get(bottomLevel, marked);
        if (iMarkedIt) {
          find(key, preds, succs);
          // size_ -= 1;
          // atomicDec(&size_, 1);
          return true;
        } else if (marked[0]) {
          return false;
        }
      }
    }
  }
}*/

__device__ bool LockFreeSkipList::Add(key_type key) {
  Node *newNode = GetNewNode(key, this->pointerIndex, calcLevel(key));
  int topLevel = newNode->topLevel;
  int bottomLevel = 0;
  Node *preds[MAX_LEVEL + 1];
  Node *succs[MAX_LEVEL + 1];
  int level;
  while (true) {
    bool found = find(key, preds, succs);
    if (found) {
      return false;
    } else {
      Node *pred;
      Node *succ;
      for (level = bottomLevel; level <= topLevel; level++) {
        succ = succs[level];
        newNode->SetRef(level, succ, false);
      }
      pred = preds[bottomLevel];
      succ = succs[bottomLevel];
      bool t;
      // printf("--- key is %d pred is %d succ is %d level is %d
      // \n",(int)key,(int)pred->key,(int)succ->key,0);
      t = pred->CompareAndSet(bottomLevel, succ, newNode, false, false);
      if (!t) {
        continue;
      }
      for (level = bottomLevel + 1; level <= topLevel; level++) {

        while (true) {
          pred = preds[level];
          succ = succs[level];
          newNode->SetRef(level, succ, false);
          // printf("-- key is %d pred is %d succ is %d level is %d
          // \n",(int)key,(int)pred->key,(int)succ->key,(int)level);
          if (pred->CompareAndSet(level, succ, newNode, false, false)) {
            break;
          }
          // printf("key is %d pred is %d succ is %d level is %d
          // \n",(int)key,(int)pred->key,(int)succ->key,(int)level);
          find(key, preds, succs);
        }
      }
      // size_ += 1;
      // this->key_map.insert(MapNode(ll, newNode));
      // atomicAdd(&size_, 1);
      return true;
    }
  }
}
#endif