/*
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */

#pragma once

#include <map>
#include <optional>
#include "Data/Concurent/MutexProtectedObject.hpp"

namespace OpenVulkano
{
	template<typename KEY, typename VALUE>
	class ThreadSafeMap
	{
		using MAP_TYPE = std::map<KEY, VALUE>;
		using SMPO_TYPE = SharedMutexProtectedObject<MAP_TYPE, SpintexShared>;
		SMPO_TYPE m_map;

	public:

		struct LockingIterator
		{
			typename MAP_TYPE::iterator iter;
			typename SMPO_TYPE::SharedAccessor accessor;

			LockingIterator(typename MAP_TYPE::iterator iter, typename SMPO_TYPE::SharedAccessor& accessor)
				: iter(iter), accessor(std::move(accessor))
			{}

			typename MAP_TYPE::reference operator*() const noexcept
			{ return *iter; }

			typename MAP_TYPE::pointer operator->() const noexcept
			{ return iter.operator->(); }

			LockingIterator& operator++() noexcept
			{
				iter++;
				return *this;
			}

			LockingIterator& operator++(int) noexcept
			{
				iter++;
				return *this;
			}

			LockingIterator& operator--() noexcept
			{
				iter--;
				return *this;
			}

			LockingIterator& operator--(int) noexcept
			{
				iter--;
				return *this;
			}

			friend bool
			operator==(const LockingIterator& lhs, const LockingIterator& rhs) noexcept
			{ return lhs.iter == rhs.iter; }

			friend bool
			operator!=(const LockingIterator& lhs, const LockingIterator& rhs) noexcept
			{ return lhs.iter != rhs.iter; }
		};

		std::optional<typename MAP_TYPE::iterator> find(KEY key)
		{
			auto access = m_map.SharedAccess();
			auto result = access->find(key);
			if (result == access->end()) return std::nullopt;
			return result;
		}

		void erase(typename MAP_TYPE::iterator iter)
		{
			auto access = m_map.Access();
			access->erase(iter);
		}

		void erase(KEY key)
		{
			auto access = m_map.Access();
			access->erase(key);
		}

		LockingIterator begin()
		{
			auto access = m_map.SharedAccess();
			return LockingIterator(access->begin(), access);
		}

		LockingIterator end()
		{
			auto access = m_map.SharedAccess();
			return LockingIterator(access->end(), access);
		}

		template<typename... _Args>
		std::pair<typename MAP_TYPE::iterator, bool> emplace(_Args&&... __args)
		{
			auto access = m_map.Access();
			return access->emplace(std::forward<_Args>(__args)...);
		}

		[[nodiscard]] size_t size() const
		{
			return m_map.UnsafeAccess().size();
		}

		void clear()
		{
			auto access = m_map.Access();
			access->clear();
		}
	};
}