#pragma once

#include <atomic.hpp>
#include <backoff.hpp>
#include <support.hpp>

#include <vector>

namespace tpcc {
namespace solutions {

class PetersonLock {
 public:
  PetersonLock() {
    want_[0].store(false);
    want_[1].store(false);
    victim_.store(0);
  }

  PetersonLock(const PetersonLock&) = delete;
  PetersonLock(PetersonLock&&) = delete;

  void Lock(const size_t thread_index) {
    want_[thread_index].store(true);
    victim_.store(thread_index);

    Backoff backoff{};
    while (want_[1 - thread_index].load() && victim_.load() == thread_index) {
      backoff();
    }
  }

  void Unlock(const size_t thread_index) {
    want_[thread_index].store(false);
  }

 private:
  tpcc::atomic<bool> want_[2];
  tpcc::atomic<size_t> victim_;
};

class TournamentTreeLock {
 public:
  explicit TournamentTreeLock(const size_t num_threads)
        : tree_size_(CountSize(num_threads)),
          locks_index_(tree_size_) {}

  TournamentTreeLock(const TournamentTreeLock&) = delete;
  TournamentTreeLock(TournamentTreeLock&&) = delete;

  void Lock(const size_t thread_index) {
    size_t current_node_index = GetThreadLeaf(thread_index);

    while (current_node_index > 0) {
      locks_index_[GetParent(current_node_index)].Lock(GetPetersonLockIndex(current_node_index));
      current_node_index = GetParent(current_node_index);
    }
  }

  void Unlock(const size_t thread_index) {
    size_t current_node_index = GetThreadLeaf(thread_index);
    std::vector<size_t> unlock_path;

    while (current_node_index > 0) {
      unlock_path.push_back(current_node_index);
      current_node_index = GetParent(current_node_index);
    }

    for (auto it = unlock_path.rbegin(); it != unlock_path.rend(); ++it) {
      locks_index_[GetParent(*it)].Unlock(GetPetersonLockIndex(*it));
    }
  }

 private:
  size_t CountSize(const size_t num_threads) const {
    size_t temp_size_ = 1;
    while (temp_size_ < num_threads) {
      temp_size_ <<= 1;
    }
    return temp_size_ - 1;
  }

  size_t GetParent(const size_t node_index) const {
    return (node_index - 1) / 2;
  }

  size_t GetThreadLeaf(const size_t thread_index) const {
    return thread_index + tree_size_;
  }

  size_t GetPetersonLockIndex(const size_t node_index) const {
    return node_index % 2;
  }

 private:
  size_t tree_size_;
  std::vector<PetersonLock> locks_index_;
};

} // namespace solutions
} // namespace tpcc