#pragma once

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

#include <twist/support/locking.hpp>

#include <functional>
#include <vector>
#include <iostream>

namespace solutions {

template <typename Key, typename Value>
class FixedSizeHashTable {
 private:
  class Bucket {
   public:
    struct Node {
      Node(const Key& key_val, const Value& val)
          : key(key_val), value(val), next(nullptr) {
      }

      twist::atomic<Node*> next;
      const Value value;
      const Key key;
    };

    bool Insert(const Key& key, const Value& value) {
      auto lock = twist::LockUnique(write_lock_);
      Node* new_node = new Node(key, value);
      if (head_.load() == nullptr) {
        head_.store(new_node);
        return true;
      } else {
        Node* temp = head_.load();
        while (temp->next.load() != nullptr) {
          if (temp->key == key) {
            delete new_node;
            return false;
          }
          temp = temp->next.load();
        }
        if (temp->key == key) {
          delete new_node;
          return false;
        }
        temp->next.store(new_node);
        return true;
      }
    }

    bool Remove(const Key& key) {
      auto lock = twist::LockUnique(write_lock_);
      if (head_.load() == nullptr) {
        return false;
      }
      if (head_.load()->key == key) {
        Node* head = head_.load();
        Node* head_next = head->next.load();
        head_.store(head_next);
        read_lock_.Synchronize();
        delete head;
        return true;
      } else {
        Node* temp = head_.load();
        while (temp->next.load() != nullptr) {
          if (temp->next.load()->key == key) {
            Node* next = temp->next.load();
            Node* new_next = next->next.load();
            temp->next.store(new_next);
            read_lock_.Synchronize();
            delete next;
            return true;
          }
          temp = temp->next.load();
        }
        return false;
      }
    }

    bool Lookup(const Key& key, Value& value) {
      auto lock = twist::LockShared(read_lock_);
      Node* temp = head_.load();
      while (temp != nullptr) {
        if (temp->key == key) {
          value = temp->value;
          return true;
        }
        temp = temp->next.load();
      }
      return false;
    }

   private:
    SpinLock write_lock_;
    RCULock read_lock_;
    twist::atomic<Node*> head_{nullptr};
  };

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

  bool Insert(const Key& key, const Value& value) {
    size_t index = (std::hash<Key>{}(key)) % buckets_.size();
    return buckets_[index].Insert(key, value);
  }

  bool Remove(const Key& key) {
    size_t index = (std::hash<Key>{}(key)) % buckets_.size();
    return buckets_[index].Remove(key);
  }

  // Wait-free
  bool Lookup(const Key& key, Value& value) {
    size_t index = (std::hash<Key>{}(key)) % buckets_.size();
    return buckets_[index].Lookup(key, value);
  }

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

}  // namespace solutions