#pragma once

#include <psemek/util/hash.hpp>
#include <psemek/util/span.hpp>
#include <psemek/util/at.hpp>

#include <memory>
#include <initializer_list>
#include <format>

namespace psemek::util
{

	namespace detail
	{

		constexpr std::uint64_t stored_value_mask = 1ull << 63;
		constexpr std::uint64_t tombstone_mask = 1ull << 62;
		constexpr std::uint64_t hash_value_mask = ~(stored_value_mask | tombstone_mask);

		template <typename T>
		struct hash_table_entry
		{
			std::uint64_t hash = 0;
			alignas(T) char storage[sizeof(T)] = {0};

			bool has_value() const
			{
				return (hash & stored_value_mask) != 0;
			}

			bool is_tombstone() const
			{
				return (hash & tombstone_mask) != 0;
			}

			T * storage_ptr()
			{
				return reinterpret_cast<T *>(storage);
			}

			T & value()
			{
				return *storage_ptr();
			}

			template <typename H>
			void set_value(H && value, std::uint64_t hash)
			{
				new (storage_ptr()) T{std::forward<H>(value)};
				this->hash = (hash & hash_value_mask) | stored_value_mask;
			}

			bool hash_equal(std::uint64_t hash) const
			{
				return (hash & hash_value_mask) == (this->hash & hash_value_mask);
			}

			void set_tombstone()
			{
				this->hash = tombstone_mask;
			}

			void reset()
			{
				if (has_value())
					value().~T();

				hash = 0;
			}

			~hash_table_entry()
			{
				reset();
			}
		};

		template <typename T>
		struct hash_table_iterator
		{
			using value_type = T;
			using pointer = T *;
			using reference = T &;
			using difference_type = std::ptrdiff_t;
			using iterator_category = std::forward_iterator_tag;

			using entry_type = hash_table_entry<std::remove_const_t<T>>;

			hash_table_iterator(entry_type * p, entry_type * end)
				: p_(p)
				, end_(end)
			{
				advance();
			}

			T & operator *() const
			{
				return p_->value();
			}

			T * operator ->() const
			{
				return std::addressof(p_->value());
			}

			hash_table_iterator<T> & operator ++()
			{
				++p_;
				advance();
				return *this;
			}

			hash_table_iterator<T> operator ++(int)
			{
				auto copy = *this;
				this->operator++();
				return copy;
			}

			friend bool operator == (hash_table_iterator<T> const & it1, hash_table_iterator<T> const & it2)
			{
				return it1.p_ == it2.p_;
			}

			hash_table_iterator<T const> as_const() const
			{
				return {p_, end_};
			}

			entry_type * internal() const
			{
				return p_;
			}

		private:
			entry_type * p_;
			entry_type * end_;

			void advance()
			{
				while (p_ != end_ && !p_->has_value())
					++p_;
			}
		};

		template <typename T>
		struct hash_table_storage
		{
			std::unique_ptr<hash_table_entry<T>[]> table;
			std::size_t capacity = 0;

			void reset(std::size_t capacity)
			{
				table.reset(new hash_table_entry<T>[capacity]);
				this->capacity = capacity;
			}

			util::span<hash_table_entry<T>> entries()
			{
				return {table.get(), table.get() + capacity};
			}

			hash_table_iterator<T> iterator(std::size_t index) const
			{
				return {table.get() + index, table.get() + capacity};
			}
		};

		template <typename T, typename Hash, typename Equal, typename KeyProjector>
		struct hash_table_impl
			: Hash, Equal, KeyProjector
		{
			hash_table_impl(Hash && hash, Equal && equal, KeyProjector && key_projector)
				: Hash(std::move(hash))
				, Equal(std::move(equal))
				, KeyProjector(std::move(key_projector))
			{}

			hash_table_impl(Hash const & hash, Equal const & equal, KeyProjector const & key_projector)
				: Hash(hash)
				, Equal(equal)
				, KeyProjector(key_projector)
			{}

			hash_table_impl(hash_table_impl && other)
				: Hash(std::move(other.hash()))
				, Equal(std::move(other.equal()))
				, storage_(std::move(other.storage_))
				, size_(other.size_)
			{
				other.storage_.capacity = 0;
				other.size_ = 0;
			}

			hash_table_impl & operator = (hash_table_impl && other)
			{
				if (this == &other)
					return *this;

				storage_ = std::move(other.storage_);
				size_ = other.size_;
				other.storage_.capacity = 0;
				other.size_ = 0;
				return *this;
			}

			Hash const & hash() const { return *this; }
			Equal const & equal() const { return *this; }
			KeyProjector const & key_projector() const { return *this; }

			template <typename H>
			std::pair<hash_table_iterator<T>, bool> insert(H && value)
			{
				// Need to make sure the key isn't present first, because
				// inserting it at a tombstone can lead to duplicate keys
				if (auto it = find(this->key_projector()(value)); it != end())
					return {it, false};

				ensure_capacity_for(size_ + 1);
				std::uint64_t hash = this->hash()(this->key_projector()(value));
				return insert_impl(std::forward<H>(value), hash);
			}

			template <typename Key>
			hash_table_iterator<T> find(Key const & key) const
			{
				if (size_ == 0)
					return end();
				std::uint64_t hash = this->hash()(key);
				return find_impl(key, hash);
			}

			void erase(hash_table_entry<T> * entry)
			{
				entry->reset();
				entry->set_tombstone();
				--size_;
				++tombstone_count_;

				// Ensure at most 25% tombstones
				if (4 * tombstone_count_ >= storage_.capacity)
					rehash();
			}

			void clear()
			{
				for (auto & entry : storage_.entries())
					entry.reset();
				size_ = 0;
				tombstone_count_ = 0;
			}

			~hash_table_impl()
			{
				clear();
			}

			hash_table_iterator<T> begin() const
			{
				return storage_.iterator(0);
			}

			hash_table_iterator<T> end() const
			{
				return storage_.iterator(storage_.capacity);
			}

			std::size_t size() const
			{
				return size_;
			}

			std::size_t capacity() const
			{
				return storage_.capacity;
			}

			void reserve(std::size_t capacity)
			{
				ensure_capacity_for(capacity);
			}

		private:
			hash_table_storage<T> storage_;
			std::size_t size_ = 0;
			std::size_t tombstone_count_ = 0;

			static std::size_t min_capacity_for_size(std::size_t size)
			{
				// Ensure at most 0.5 load factor
				return 2 * size;
			}

			static std::size_t find_capacity(std::size_t current_capacity, std::size_t min_capacity)
			{
				current_capacity = std::max(current_capacity, std::size_t(16));
				while (current_capacity < min_capacity)
					current_capacity *= 2;
				return current_capacity;
			}

			void ensure_capacity_for(std::size_t size)
			{
				std::size_t capacity = min_capacity_for_size(size);
				if (storage_.capacity < capacity)
					reallocate(find_capacity(storage_.capacity, capacity));
			}

			void reallocate(std::size_t capacity)
			{
				hash_table_storage<T> storage;
				storage.reset(capacity);

				std::swap(storage_, storage);

				size_ = 0;
				tombstone_count_ = 0;

				for (hash_table_entry<T> & entry : storage.entries())
				{
					if (entry.has_value())
					{
						insert_impl(std::move(entry.value()), entry.hash);
						entry.reset();
					}
				}
			}

			void rehash()
			{
				reallocate(capacity());
			}

			std::size_t probe_index(std::uint64_t hash, std::size_t i) const
			{
				return (static_cast<std::size_t>(hash) + (i * (i + 1)) / 2) % storage_.capacity;
			}

			template <typename H>
			std::pair<hash_table_iterator<T>, bool> insert_impl(H && value, std::uint64_t hash)
			{
				std::size_t i = 0;
				while (true)
				{
					std::size_t index = probe_index(hash, i);
					auto & entry = storage_.table[index];
					if (!entry.has_value() || entry.is_tombstone())
					{
						entry.set_value(std::forward<H>(value), hash);
						++size_;
						return {storage_.iterator(index), true};
					}
					else if (entry.hash_equal(hash) && equal()(key_projector()(value), key_projector()(entry.value())))
					{
						return {storage_.iterator(index), false};
					}
					else
						++i;
				}
			}

			template <typename Key>
			hash_table_iterator<T> find_impl(Key const & key, std::uint64_t hash) const
			{
				std::size_t i = 0;
				while (true)
				{
					std::size_t index = probe_index(hash, i);
					auto & entry = storage_.table[index];
					if (!entry.is_tombstone())
					{
						if (!entry.has_value())
						{
							return end();
						}
						else if (entry.hash_equal(hash) && equal()(key, key_projector()(entry.value())))
						{
							return storage_.iterator(index);
						}
					}
					++i;
				}
			}
		};

		struct id_key_projector
		{
			template <typename Key>
			Key const & operator() (Key const & key) const
			{
				return key;
			}
		};

		struct pair_key_projector
		{
			template <typename Key, typename Value>
			Key const & operator() (std::pair<Key, Value> const & pair) const
			{
				return pair.first;
			}
		};

	}

	template <typename T, typename Hash = std::hash<T>, typename Equal = std::equal_to<T>>
	struct hash_set
	{
		using iterator = detail::hash_table_iterator<T const>;

		hash_set(Hash && hash = {}, Equal && equal = {})
			: impl_(hash, equal, {})
		{}

		hash_set(std::initializer_list<T> init, Hash && hash = {}, Equal && equal = {})
			: impl_(hash, equal, {})
		{
			for (auto & value : init)
				insert(std::move(value));
		}

		hash_set(hash_set && other) = default;

		hash_set(hash_set const & other) requires std::is_copy_constructible_v<T>
			: impl_(other.impl_.hash(), other.impl_.equal(), other.impl_.key_projector())
		{
			for (auto const & value : other)
				insert(value);
		}

		hash_set & operator = (hash_set && other) = default;

		hash_set & operator = (hash_set const & other) requires std::is_copy_constructible_v<T>
		{
			if (this != &other)
			{
				clear();
				for (auto const & value : other)
					insert(value);
			}
			return *this;
		}

		std::pair<iterator, bool> insert(T const & value)
		{
			auto result = impl_.insert(value);
			return {result.first.as_const(), result.second};
		}

		std::pair<iterator, bool> insert(T && value)
		{
			auto result = impl_.insert(std::move(value));
			return {result.first.as_const(), result.second};
		}

		iterator find(T const & key) const
		{
			return impl_.find(key).as_const();
		}

		template <typename Key>
		iterator find(Key const & key) const
		{
			return impl_.find(key).as_const();
		}

		bool contains(T const & key) const
		{
			return find(key) != end();
		}

		template <typename Key>
		bool contains(Key const & key) const
		{
			return find(key) != end();
		}

		bool erase(iterator const & it)
		{
			impl_.erase(it.internal());
			return true;
		}

		bool erase(T const & key)
		{
			if (auto it = find(key); it != end())
			{
				erase(it);
				return true;
			}
			return false;
		}

		template <typename Key>
		bool erase(Key const & key)
		{
			if (auto it = find(key); it != end())
			{
				erase(it);
				return true;
			}
			return false;
		}

		iterator begin() const
		{
			return impl_.begin().as_const();
		}

		iterator end() const
		{
			return impl_.end().as_const();
		}

		void clear()
		{
			impl_.clear();
		}

		bool empty() const
		{
			return impl_.size() == 0;
		}

		std::size_t size() const
		{
			return impl_.size();
		}

		void reserve(std::size_t capacity)
		{
			impl_.reserve(capacity);
		}

	private:
		detail::hash_table_impl<T, Hash, Equal, detail::id_key_projector> impl_;
	};

	template <typename Key, typename Value, typename KeyHash = std::hash<Key>, typename KeyEqual = std::equal_to<Key>>
	struct hash_map
	{
		using iterator = detail::hash_table_iterator<std::pair<Key const, Value>>;

		hash_map(KeyHash && hash = {}, KeyEqual && equal = {})
			: impl_(hash, equal, {})
		{}

		hash_map(std::initializer_list<std::pair<Key, Value>> init, KeyHash && hash = {}, KeyEqual && equal = {})
			: impl_(hash, equal, {})
		{
			for (auto & pair : init)
				insert(std::move(pair));
		}

		hash_map(hash_map && other) = default;

		hash_map(hash_map const & other) requires std::is_copy_constructible_v<std::pair<Key, Value>>
			: impl_(other.impl_.hash(), other.impl_.equal(), other.impl_.key_projector())
		{
			for (auto const & pair : other)
				insert({pair.first, pair.second});
		}

		hash_map & operator = (hash_map && other) = default;

		hash_map & operator = (hash_map const & other) requires std::is_copy_constructible_v<std::pair<Key, Value>>
		{
			if (this != &other)
			{
				clear();
				for (auto const & pair : other)
					insert({pair.first, pair.second});
			}
			return *this;
		}

		std::pair<iterator, bool> insert(std::pair<Key, Value> const & value)
		{
			return impl_.insert(value);
		}

		std::pair<iterator, bool> insert(std::pair<Key, Value> && value)
		{
			return impl_.insert(std::move(value));
		}

		iterator find(Key const & key) const
		{
			return impl_.find(key);
		}

		template <typename Key1>
		iterator find(Key1 const & key) const
		{
			return impl_.find(key);
		}

		bool contains(Key const & key) const
		{
			return find(key) != end();
		}

		template <typename Key1>
		bool contains(Key1 const & key) const
		{
			return find(key) != end();
		}

		bool erase(iterator const & it)
		{
			impl_.erase(it.internal());
			return true;
		}

		bool erase(Key const & key)
		{
			if (auto it = find(key); it != end())
			{
				erase(it);
				return true;
			}
			return false;
		}

		template <typename Key1>
		bool erase(Key1 const & key)
		{
			if (auto it = find(key); it != end())
			{
				erase(it);
				return true;
			}
			return false;
		}

		iterator begin() const
		{
			return impl_.begin();
		}

		iterator end() const
		{
			return impl_.end();
		}

		Value & operator[] (Key const & key)
		{
			if (auto it = find(key); it != end())
				return it->second;
			return insert({Key(key), Value{}}).first->second;
		}

		template <typename Key1>
		Value & operator[] (Key1 const & key)
		{
			if (auto it = find(key); it != end())
				return it->second;
			return insert({Key(key), Value{}}).first->second;
		}

		Value & at(Key const & key)
		{
			auto it = find(key);
			if (it == end())
				throw util::key_error<Key>(key);
			return it->second;
		}

		template <typename Key1>
		Value & at(Key1 const & key)
		{
			auto it = find(key);
			if (it == end())
				throw util::key_error<Key>(Key(key));
			return it->second;
		}

		Value const & at(Key const & key) const
		{
			auto it = find(key);
			if (it == end())
				throw util::key_error<Key>(key);
			return it->second;
		}

		template <typename Key1>
		Value const & at(Key1 const & key) const
		{
			auto it = find(key);
			if (it == end())
				throw util::key_error<Key>(Key(key));
			return it->second;
		}

		void clear()
		{
			impl_.clear();
		}

		std::size_t size() const
		{
			return impl_.size();
		}

		bool empty() const
		{
			return impl_.size() == 0;
		}

		void reserve(std::size_t capacity)
		{
			impl_.reserve(capacity);
		}

	private:
		detail::hash_table_impl<std::pair<Key const, Value>, KeyHash, KeyEqual, detail::pair_key_projector> impl_;
	};

}


namespace std
{

	template <typename T, typename Hash, typename Equal, typename Char>
	struct formatter<::psemek::util::hash_set<T, Hash, Equal>, Char>
		: std::formatter<T, Char>
	{
		using std::formatter<T, Char>::parse;

		template <typename FormatContext>
		auto format(::psemek::util::hash_set<T, Hash, Equal> const & s, FormatContext & ctx) const
		{
			bool first = true;
			ctx.advance_to(std::format_to(ctx.out(), "{{"));
			for (auto const & v : s)
			{
				if (!first)
					ctx.advance_to(std::format_to(ctx.out(), ", "));
				first = false;
				ctx.advance_to(formatter<T, Char>::format(v, ctx));
			}
			return std::format_to(ctx.out(), "}}");
		}
	};

	template <typename Key, typename Value, typename KeyHash, typename KeyEqual, typename Char>
	struct formatter<::psemek::util::hash_map<Key, Value, KeyHash, KeyEqual>, Char>
		: std::pair<formatter<Key, Char>, formatter<Value, Char>>
	{
		template <typename ParseContext>
		constexpr auto parse(ParseContext & ctx)
		{
			return std::pair<formatter<Key, Char>, formatter<Value, Char>>::first.parse(ctx);
		}

		template <typename FormatContext>
		auto format(::psemek::util::hash_map<Key, Value, KeyHash, KeyEqual> const & m, FormatContext & ctx) const
		{
			bool first = true;
			ctx.advance_to(std::format_to(ctx.out(), "{{"));
			for (auto const & p : m)
			{
				if (!first)
					ctx.advance_to(std::format_to(ctx.out(), ", "));
				first = false;
				ctx.advance_to(std::pair<formatter<Key, Char>, formatter<Value, Char>>::first.format(p.first, ctx));
				ctx.advance_to(std::format_to(ctx.out(), ": "));
				ctx.advance_to(std::pair<formatter<Key, Char>, formatter<Value, Char>>::second.format(p.second, ctx));
			}
			return std::format_to(ctx.out(), "}}");
		}
	};

}
