#pragma once

#include <psemek/util/functional.hpp>
#include <psemek/util/hash_table.hpp>
#include <psemek/util/at.hpp>

#include <list>

namespace psemek::util
{

	// No function except `insert` touches the accessed elements
	// To mark an element as accessed, one must call `touch` manually
	// Calling touch from inside the `removable` predicate or while iterating is UB
	template <typename Key, typename Mapped, typename Predicate = decltype(always_true)>
	struct lru_cache
	{
		using key_type = Key;
		using mapped_type = Mapped;
		using value_type = std::pair<Key const, Mapped>;

		using iterator = typename std::list<value_type>::iterator;
		using const_iterator = typename std::list<value_type>::const_iterator;

		// Predicate must be callable with arguments of type (Key const, Value)
		lru_cache(std::size_t max_size, Predicate removable = Predicate{});

		// Find an element without touching it
		iterator find(Key const & key);
		const_iterator find(Key const & key) const;

		bool contains(Key const & key) const;

		// Access an element without touching it
		// Throws if the key is not present
		Mapped & at(Key const & key);
		Mapped const & at(Key const & key) const;

		// Insert an element (automatically touches it)
		// N.B. the element might be immediately deleted by a shrink
		void insert(Key const & key, Mapped && mapped);
		void insert(Key const & key, Mapped const & mapped);

		void erase(const_iterator it);
		void erase(Key const & key);

		void touch(const_iterator it);
		void touch(Key const & key);

		iterator begin();
		iterator end();

		const_iterator begin() const;
		const_iterator end() const;

		const_iterator cbegin() const;
		const_iterator cend() const;

		bool empty() const;

		std::size_t size() const;

		std::size_t max_size() const;
		std::size_t max_size(std::size_t new_max_size);

		void shrink();

		Predicate & removable();
		Predicate const & removable() const;

		bool removable(const_iterator it) const;
		bool removable(Key const & key, Mapped const & value) const;

	private:
		std::size_t max_size_;
		Predicate removable_;

		// N.B.: keys are stored twice - one in map_, one in list_
		// list_ contains non-removable items first, then all removable items
		std::list<value_type> list_;
		iterator removable_begin_ = list_.end();
		util::hash_map<Key, iterator> map_;

		iterator find_safe(Key const & key);
		const_iterator find_safe(Key const & key) const;

		void insert(iterator it);
	};

	template <typename Key, typename Mapped, typename Predicate>
	lru_cache<Key, Mapped, Predicate>::lru_cache(std::size_t max_size, Predicate removable)
		: max_size_(max_size)
		, removable_(std::move(removable))
	{}

	template <typename Key, typename Mapped, typename Predicate>
	auto lru_cache<Key, Mapped, Predicate>::find(Key const & key) -> iterator
	{
		if (auto it = map_.find(key); it != map_.end())
			return it->second;
		return list_.end();
	}

	template <typename Key, typename Mapped, typename Predicate>
	auto lru_cache<Key, Mapped, Predicate>::find(Key const & key) const -> const_iterator
	{
		return const_cast<lru_cache &>(*this).find(key);
	}

	template <typename Key, typename Mapped, typename Predicate>
	bool lru_cache<Key, Mapped, Predicate>::contains(Key const & key) const
	{
		return map_.contains(key);
	}

	template <typename Key, typename Mapped, typename Predicate>
	auto lru_cache<Key, Mapped, Predicate>::at(Key const & key) -> Mapped &
	{
		return find_safe(key)->second;
	}

	template <typename Key, typename Mapped, typename Predicate>
	auto lru_cache<Key, Mapped, Predicate>::at(Key const & key) const -> Mapped const &
	{
		return const_cast<lru_cache &>(*this).at(key);
	}

	template <typename Key, typename Mapped, typename Predicate>
	void lru_cache<Key, Mapped, Predicate>::insert(Key const & key, Mapped && mapped)
	{
		list_.emplace_front(key, std::move(mapped));
		insert(list_.begin());
	}

	template <typename Key, typename Mapped, typename Predicate>
	void lru_cache<Key, Mapped, Predicate>::insert(Key const & key, Mapped const & mapped)
	{
		list_.emplace_front(key, mapped);
		insert(list_.begin());
	}

	template <typename Key, typename Mapped, typename Predicate>
	void lru_cache<Key, Mapped, Predicate>::erase(const_iterator it)
	{
		if (it == removable_begin_)
			++removable_begin_;
		map_.erase(it->first);
		list_.erase(it);
	}

	template <typename Key, typename Mapped, typename Predicate>
	void lru_cache<Key, Mapped, Predicate>::erase(Key const & key)
	{
		erase(find_safe(key));
	}

	template <typename Key, typename Mapped, typename Predicate>
	void lru_cache<Key, Mapped, Predicate>::touch(const_iterator it)
	{
		bool const removable = removable_(it->first, it->second);
		list_.splice(removable ? removable_begin_ : list_.begin(), list_, it);
		if (removable && it != removable_begin_)
			--removable_begin_;
	}

	template <typename Key, typename Mapped, typename Predicate>
	void lru_cache<Key, Mapped, Predicate>::touch(Key const & key)
	{
		touch(find_safe(key));
	}

	template <typename Key, typename Mapped, typename Predicate>
	auto lru_cache<Key, Mapped, Predicate>::begin() -> iterator
	{
		return list_.begin();
	}

	template <typename Key, typename Mapped, typename Predicate>
	auto lru_cache<Key, Mapped, Predicate>::end() -> iterator
	{
		return list_.end();
	}

	template <typename Key, typename Mapped, typename Predicate>
	auto lru_cache<Key, Mapped, Predicate>::begin() const -> const_iterator
	{
		return list_.begin();
	}

	template <typename Key, typename Mapped, typename Predicate>
	auto lru_cache<Key, Mapped, Predicate>::end() const -> const_iterator
	{
		return list_.end();
	}

	template <typename Key, typename Mapped, typename Predicate>
	auto lru_cache<Key, Mapped, Predicate>::cbegin() const -> const_iterator
	{
		return list_.begin();
	}

	template <typename Key, typename Mapped, typename Predicate>
	auto lru_cache<Key, Mapped, Predicate>::cend() const -> const_iterator
	{
		return list_.end();
	}

	template <typename Key, typename Mapped, typename Predicate>
	bool lru_cache<Key, Mapped, Predicate>::empty() const
	{
		return list_.empty();
	}

	template <typename Key, typename Mapped, typename Predicate>
	std::size_t lru_cache<Key, Mapped, Predicate>::size() const
	{
		return list_.size();
	}

	template <typename Key, typename Mapped, typename Predicate>
	std::size_t lru_cache<Key, Mapped, Predicate>::max_size() const
	{
		return max_size_;
	}

	template <typename Key, typename Mapped, typename Predicate>
	std::size_t lru_cache<Key, Mapped, Predicate>::max_size(std::size_t new_max_size)
	{
		std::swap(max_size_, new_max_size);
		shrink();
		return new_max_size;
	}

	template <typename Key, typename Mapped, typename Predicate>
	void lru_cache<Key, Mapped, Predicate>::shrink()
	{
		while (size() > max_size())
		{
			if (removable_begin_ != list_.end())
			{
				bool const removable = removable_(list_.back().first, list_.back().second);
				if (removable)
				{
					erase(std::prev(list_.end()));
				}
				else
				{
					// Move the item to non-removable sublist
					// We will have one potentially removable item less in the next loop iteration
					if (removable_begin_ == std::prev(list_.end()))
						++removable_begin_;
					else
						list_.splice(removable_begin_, list_, std::prev(list_.end()));
				}
			}
			else
			{
				// No potentially removable elements, try non-removable ones
				if (list_.begin() == removable_begin_)
					break;

				auto it = std::prev(removable_begin_);
				if (!removable_(it->first, it->second))
					break;

				erase(it);
			}
		}
	}

	template <typename Key, typename Mapped, typename Predicate>
	Predicate & lru_cache<Key, Mapped, Predicate>::removable()
	{
		return removable_;
	}

	template <typename Key, typename Mapped, typename Predicate>
	Predicate const & lru_cache<Key, Mapped, Predicate>::removable() const
	{
		return removable_;
	}

	template <typename Key, typename Mapped, typename Predicate>
	bool lru_cache<Key, Mapped, Predicate>::removable(const_iterator it) const
	{
		return removable_(it->first, it->second);
	}

	template <typename Key, typename Mapped, typename Predicate>
	bool lru_cache<Key, Mapped, Predicate>::removable(Key const & key, Mapped const & value) const
	{
		return removable_(key, value);
	}

	template <typename Key, typename Mapped, typename Predicate>
	auto lru_cache<Key, Mapped, Predicate>::find_safe(Key const & key) -> iterator
	{
		if (auto it = find(key); it != list_.end())
			return it;
		throw key_error{key};
	}

	template <typename Key, typename Mapped, typename Predicate>
	auto lru_cache<Key, Mapped, Predicate>::find_safe(Key const & key) const -> const_iterator
	{
		return const_cast<lru_cache &>(*this).find_safe(key);
	}

	template <typename Key, typename Mapped, typename Predicate>
	void lru_cache<Key, Mapped, Predicate>::insert(iterator it)
	{
		bool const removable = removable_(it->first, it->second);
		if (removable)
		{
			list_.splice(removable_begin_, list_, it);
			it = --removable_begin_;
		}
		map_[it->first] = it;
		shrink();
	}

}
