/*
    libmaus2
    Copyright (C) 2020 German Tischler-Höhle

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/
#if ! defined(LOCKEDMAP_HPP)
#define LOCKEDMAP_HPP

#include <memory>
#include <mutex>
#include <libmaus2/util/atomic_shared_ptr.hpp>
#include <libmaus2/avl/AtomicAVLPtrValueMap.hpp>

namespace libmaus2
{
	namespace parallel
	{
		namespace threadpool
		{
			template<typename value_type>
			struct LockedMapDefaultAllocator
			{
				libmaus2::util::shared_ptr<value_type> allocate() const
				{
					libmaus2::util::shared_ptr<value_type> ptr(new value_type);
					return ptr;
				}
			};

			template<typename _value_type>
			struct AtomicAllocator
			{
				typedef _value_type value_type;
				typedef AtomicAllocator<value_type> this_type;
				typedef std::unique_ptr<this_type> unique_ptr_type;
				typedef std::shared_ptr<this_type> shared_ptr_type;

				libmaus2::util::shared_ptr< std::atomic<value_type> > allocate() const
				{
					libmaus2::util::shared_ptr< std::atomic<value_type> > tptr(
						new std::atomic<value_type>(value_type())
					);
					return tptr;
				}
			};

			template<typename _key_type, typename _value_type, typename _allocator_type = LockedMapDefaultAllocator<_value_type> >
			struct LockedMap
			{
				typedef _key_type key_type;
				typedef _value_type value_type;
				typedef _allocator_type allocator_type;
				typedef LockedMap<key_type,value_type,allocator_type> this_type;
				typedef std::unique_ptr<this_type> unique_ptr_type;
				typedef std::shared_ptr<this_type> shared_ptr_type;

				private:
				this_type & operator=(this_type const & O) = delete;

				public:
				std::mutex lock;
				libmaus2::avl::AtomicAVLPtrValueMap<key_type,value_type> M;
				libmaus2::util::shared_ptr<allocator_type> allocator;

				LockedMap(libmaus2::util::shared_ptr<allocator_type> rallocator = libmaus2::util::shared_ptr<allocator_type>(new allocator_type)) : lock(), M(), allocator(rallocator)
				{
				}

				libmaus2::util::shared_ptr<value_type> get(key_type const & key)
				{
					std::lock_guard<std::mutex> slock(lock);
					auto it = M.find(key);

					if ( it == M.end() )
					{
						libmaus2::util::shared_ptr<value_type> ptr(allocator->allocate());
						M.insert(key,ptr);
						it = M.find(key);
					}

					assert ( it != M.end() );

					return it->second.load();
				}
			};
		}
	}
}
#endif
