#pragma once

#include <cstdlib>
#include <sys/mman.h>
#include <tpcc/lockfree/x86_64_atomic_stamped_pointer.hpp>
#include <tpcc/thread/thread_local.hpp>
#include <type_traits>

namespace tpcc {
  namespace detail {
    template <size_t kChunkSize, class=std::enable_if_t<(kChunkSize >= sizeof(void*))>>
    class FixedSizeAllocator {
     private:
      struct Node {
        std::atomic<Node*> next_;

      };

      struct Chunk {
        void*& AsPointer() const {
          return *static_cast<void**>(data_);
        }

        Node* AsNode() const {
          return static_cast<Node*>(static_cast<void*>(data_));
        }

        uint8_t data_[kChunkSize];
      };

      struct Arena {
        static const size_t kChunksInArena = 1024;
        static const size_t kArenaSize = kChunkSize * kChunksInArena;
        Chunk chunks_[kChunksInArena];
      };

      struct ArenaStack {
        ArenaStack() {
        }

        ~ArenaStack() {
          while (lastArena_ != nullptr) {
            Arena* nextArena = static_cast<Arena*>(lastArena_->chunks_[0].AsPointer());
            munmap(lastArena_, Arena::kArenaSize);
            lastArena_ = nextArena;
          }
        }

        Chunk* NewChunk() {
          if (freeChunksInLastArena_ == 0) {
            AllocateNewArena();
          }
          return &lastArena_->chunks_[freeChunksInLastArena_--];
        }

        void AllocateNewArena() {
          Arena* newArena = static_cast<Arena*>(mmap(NULL, Arena::kArenaSize, PROT_WRITE, MAP_ANONYMOUS | MAP_SHARED, -1, 0));
          newArena->chunks_[0].AsPointer() = static_cast<void*>(lastArena_);
          lastArena_ = newArena;
          freeChunksInLastArena_ = Arena::kChunksInArena - 1;
        }

        size_t freeChunksInLastArena_{0};
        Arena* lastArena_{nullptr};
      };

      struct NodeStack {
        using NodePtr = tpcc::lockfree::X86AtomicStampedPointer<Node>;

        NodeStack() {
        }

        ~NodeStack() {
        }

        void PushNode(Node* node) {
          typename NodePtr::StampedPtr oldTop;
          do {
            oldTop = top_.Load();
            node->next_.store(oldTop.raw_ptr_);
          } while (!top_.CompareAndSet(oldTop, node));
        }

        Node* PopNode() {
          typename NodePtr::StampedPtr oldTop;
          Node* newTop;
          do {
            oldTop = top_.Load();
            if (!oldTop) {
              return nullptr;
            }
            newTop = oldTop->next_.load();
          } while (!top_.CompareAndSet(oldTop, newTop));
          return oldTop.raw_ptr_;
        }

        NodePtr top_{nullptr};
      };

     public:
      FixedSizeAllocator() {
      }

      ~FixedSizeAllocator() {
      }

      void* New() {
        Node* node = pool_.PopNode();
        if (node == nullptr) {
          return arenaList_.Get().NewChunk();
        } else {
          return node;
        }
      }

      void Delete(void* ptr) {
        Node* node = static_cast<Node*>(ptr);
        pool_.PushNode(node);
      }

     private:
      NodeStack pool_;
      tpcc::thread::ThreadLocal<ArenaStack> arenaList_;
    };
  } // namespace detail
} // namespace tpcc