#pragma once

#include "rcu_lock.hpp"
#include "spinlock.hpp"

#include <ctime>
#include <twist/support/locking.hpp>

#include <functional>
#include <vector>

namespace solutions {

template <typename Key, typename Value>
class FixedSizeHashTable {
 private:
  class ListNode {
   public:
    ListNode(Key key, Value value, ListNode* next)
        : key_(key), value_(value), next_(next) {
    }

    Key GetKey() {
      return key_;
    }

    Value GetValue() {
      return value_;
    }

    ListNode* GetNext() {
      return next_.load();
    }

    void SetNext(ListNode* next) {
      next_.store(next);
    }

   private:
    Key key_;
    Value value_;
    twist::atomic<ListNode*> next_;
  };

  class Bucket {
   public:
    ListNode* GetHead() {
      return head_.load();
    }

    void SetHead(ListNode* new_head) {
      head_.store(new_head);
    }

    solutions::SpinLock* GetSpinLock() {
      return &spinlock_;
    }

    solutions::RCULock* GetRCULock() {
      return &rculock_;
    }

   private:
    solutions::RCULock rculock_;
    solutions::SpinLock spinlock_;
    twist::atomic<ListNode*> head_{nullptr};
  };

 public:
  FixedSizeHashTable(size_t num_buckets) : buckets_(num_buckets) {
  }

  bool Insert(const Key& key, const Value& value) {
    auto bucket = GetBucket(key);
    solutions::SpinLock* spinlock = bucket->GetSpinLock();
    auto lock = twist::LockUnique<solutions::SpinLock>(*spinlock);

    auto curr = bucket->GetHead();
    while (curr != nullptr) {
      if (curr->GetKey() == key) {
        return false;
      }
      curr = curr->GetNext();
    }

    auto new_node = new ListNode(key, value, bucket->GetHead());
    bucket->SetHead(new_node);
    return true;
  }

  bool Remove(const Key& key) {
    auto bucket = GetBucket(key);
    solutions::SpinLock* spinlock = bucket->GetSpinLock();
    auto lock = twist::LockUnique<solutions::SpinLock>(*spinlock);

    ListNode* curr = nullptr;
    ListNode* next = bucket->GetHead();

    if (next == nullptr) {
      return false;
    }

    if (next != nullptr && next->GetKey() == key) {
      bucket->SetHead(next->GetNext());
      solutions::RCULock* rculock = bucket->GetRCULock();
      rculock->Synchronize();
      delete (ListNode*)next;
      return true;
    }

    curr = next;
    next = curr->GetNext();
    while (next != nullptr) {
      if (next->GetKey() == key) {
        curr->SetNext(next->GetNext());
        solutions::RCULock* rculock = bucket->GetRCULock();
        rculock->Synchronize();
        delete (ListNode*)next;
        return true;
      }
      curr = next;
      next = curr->GetNext();
    }

    return false;
  }

  // Wait-free
  bool Lookup(const Key& key, Value& value) {
    Bucket* bucket = GetBucket(key);
    solutions::RCULock* rculock = bucket->GetRCULock();
    auto lock = twist::LockShared<solutions::RCULock>(*rculock);

    auto curr = bucket->GetHead();
    while (curr != nullptr) {
      if (curr->GetKey() == key) {
        value = curr->GetValue();
        return true;
      }
      curr = curr->GetNext();
    }
    return false;
  }

 private:
  Bucket* GetBucket(const Key& key) {
    return &buckets_[std::hash<Key>()(key) % buckets_.size()];
  }

 private:
  std::vector<Bucket> buckets_;
};

}  // namespace solutions