#pragma once

#include <stdexcept>
#include <vector>
#include <stack>
#include <memory>

using namespace std;


template <typename T, typename Index>
class ListAlloc {
	static_assert(is_integral_v<Index> && is_unsigned_v<Index>);

public:
	const size_t capacity;

	ListAlloc(size_t capacity);
	~ListAlloc();

	template <typename... Args>
	Index allocate(Args... args);
	void deallocate(Index index);

	T* at(Index index);

private:
	char *buffer;
	stack<Index> freeStack;
	size_t cursor = 0;
};

template <typename T, typename Index>
class ListAllocRef {
	static_assert(is_integral_v<Index> && is_unsigned_v<Index>);

public:
	ListAllocRef() : index(-1) {}
	ListAllocRef(Index index) : index(index) {}

	~ListAllocRef() {
		if (index != (Index)-1) {
			assert(false && "Non-empty ListAllocRef upon destruction");
		}
	}

	// No copying
	ListAllocRef(const ListAllocRef<T, Index> &other) = delete;
	ListAllocRef<T, Index>& operator=(const ListAllocRef<T, Index> &other) = delete;

	operator bool() const {
		return index != (Index)-1;
	}

	ListAllocRef<T, Index>& operator=(Index newIndex) {
		if (index != (Index)-1) {
			throw logic_error("operator= on non-empty ListAllocRef");
		}

		index = newIndex;
		return *this;
	}

	T* get(ListAlloc<T, Index> &allocator) {
		return allocator.at(index);
	}

	const T* get(ListAlloc<T, Index> &allocator) const {
		return allocator.at(index);
	}

	void deallocate(ListAlloc<T, Index> &allocator) {
		if (index != (Index)-1) {
			allocator.deallocate(index);
			index = -1;
		}
	}

private:
	Index index;
};

template <typename T, typename Index>
ListAlloc<T, Index>::ListAlloc(size_t capacity)
		: capacity(capacity)
		, buffer(new char[capacity * sizeof(T)]) {

	size_t largestIndex = capacity - 1;
	if (capacity != 0 && (size_t)(Index)largestIndex != largestIndex) {
		throw logic_error("Capacity too large for index type in ListAlloc");
	}

	fprintf(stderr, "ListAlloc with capacity=%zu, size=%zu\n", capacity, capacity * sizeof(T));
}

template <typename T, typename Index>
ListAlloc<T, Index>::~ListAlloc() {
	if (freeStack.size() != cursor) {
		assert(false && "Not all entries deallocated in ~ListAlloc");
	}

	delete[] buffer;
}

template <typename T, typename Index>
template <typename... Args>
Index ListAlloc<T, Index>::allocate(Args... args) {
	Index index;

	if (!freeStack.empty()) {
		index = freeStack.top();
		freeStack.pop();
	} else if (cursor < capacity) {
		index = cursor++;
	} else {
		throw runtime_error("Out of memory in ListAlloc");
	}

	new(at(index)) T(args...);
	return index;
}

template <typename T, typename Index>
void ListAlloc<T, Index>::deallocate(Index index) {
	at(index)->~T();
	freeStack.push(index);
}

template <typename T, typename Index>
T* ListAlloc<T, Index>::at(Index index) {
	return (T*)&buffer[index * sizeof(T)];
}