#include <cstddef>
#include <stdexcept>
#include <iostream>
#include <bitset>
#include <memory>

class Arena
{
	uint8_t* allocation_map;
	uint8_t* data;
	std::size_t arena_size;
	std::size_t block_size;
	std::size_t n_blocks;

	std::size_t hole_offset = 0;
	std::size_t hole_offset_size;

	public:
		Arena(Arena const &) = delete;
		Arena& operator=(Arena const&) = delete;

		Arena(std::size_t n_blocks, std::size_t block_size)
			: allocation_map(new uint8_t[n_blocks]()), data(new uint8_t[n_blocks * block_size]()),
			arena_size(n_blocks * block_size), hole_offset_size(arena_size), block_size(block_size), n_blocks(n_blocks)
		{
#ifdef DEBUG
			std::cout << "Constructing Arena [" << this << "] of " << this->arena_size << " bytes." << std::endl;
#endif
		}

		~Arena()
		{
#ifdef DEBUG
			std::cout << "Destructing Arena [" << this << "] of " << this->arena_size << " bytes." << std::endl;
#endif
			delete[] allocation_map;
			delete[] data;
		}

		void* allocate(std::size_t n)
		{
			// Allocate n blocks;
#ifdef DEBUG
			std::cout << "Allocating " << n << " blocks." << std::endl;
			std::cout << "== Allocation map at start of allocation: " << std::endl;
			this->print_allocation_map();
#endif

			int contiguous_size = 0;
			int index = hole_offset / block_size;

			while ((contiguous_size < n) && (index < n_blocks))
			{
				if ((bool)allocation_map[index])
					contiguous_size = 0;
				else
					contiguous_size++;
				index++;
			}

#ifdef DEBUG
			std::cout << "Contiguous size: " << contiguous_size << std::endl;
			std::cout << "Index: " << index << std::endl;
			std::cout << "Requested number of blocks: " << n << std::endl;
#endif

			if (hole_offset_size >= n * block_size)
			{
				hole_offset += n * block_size;
				hole_offset_size -= n * block_size;
			}

			while (contiguous_size--)
				allocation_map[(index - n)+contiguous_size] = 1;

#ifdef DEBUG
			std::cout << "== Allocation map at end of allocation: " << std::endl;
			this->print_allocation_map();
#endif
			return this->data + (block_size * (index-n));
		}

		void deallocate(void *ptr, std::size_t n)
		{
			// Deallocate n blocks from *ptr;
			std::size_t offset = ((uint8_t*)ptr - this->data);
			std::size_t block_index = offset / block_size;
			std::size_t size = n * block_size;
#ifdef DEBUG
			this->print_allocation_map();
			std::cout << "Deallocating " << n << " blocks." << std::endl;
			std::cout << "Offset into the array: " << offset << std::endl;
#endif

			if (hole_offset > offset)
			{
				/* If the hole offset is further into the array, reset it to here. */
				hole_offset = offset;
				hole_offset_size = size;
			}

			for(int i = 0; i < n; i++)
			{
				allocation_map[block_index + i] = 0;
			}
			this->print_allocation_map();
		}

		void print_allocation_map()
		{
#ifdef DEBUG
			std::cout << "Allocation map (addr: " << (void*)allocation_map << ", block size: " << block_size << ", hole offset:" << hole_offset / block_size << ")" << std::endl;
			for (int i = 0; i < n_blocks; i++) 
			{
				std::cout << (int)allocation_map[i];
			}
			std::cout << std::endl;
#endif
		}
};

template<typename T, std::size_t N>
struct ArenaAllocator
{
	template <typename U, std::size_t S> friend struct ArenaAllocator;

	using value_type = T;

	using propagate_on_container_copy_assignment = std::true_type;
	using propagate_on_container_move_assignment = std::true_type;
	using propagate_on_container_swap = std::true_type;
	using is_always_equal = std::false_type;

	template<typename U>
	struct rebind { using other = ArenaAllocator<U, N>; };

	ArenaAllocator() : arena(new Arena(N, sizeof(T))) {};

	template<typename U>
	ArenaAllocator(ArenaAllocator<U, N> const &rhs) : arena(rhs.arena)
	{
		std::cout << "Copy constructor called with sizeof U = " << sizeof(U) << std::endl;
	}

	ArenaAllocator& operator=(const ArenaAllocator& rhs) = delete;

	value_type* allocate(std::size_t n)
	{
		return static_cast<value_type*>(arena->allocate(n));
	}

	void deallocate(value_type* p, std::size_t n)
	{
		arena->deallocate(p, n);
	}

	template<typename U>
	bool operator==(ArenaAllocator<U, N> const& rhs) const
	{
		return arena == rhs.arena;
	}

	private:
	std::shared_ptr<Arena> arena;
};

