/*
    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(LIBMAUS2_PARALLEL_THREADPOOL_HPP)
#define LIBMAUS2_PARALLEL_THREADPOOL_HPP

#include <ostream>
#include <atomic>
#include <memory>
#include <typeinfo>
#include <typeindex>
#include <mutex>
#include <stack>
#include <unordered_map>
#include <condition_variable>
#include <queue>
#include <thread>
#include <cassert>

#include <libmaus2/demangle/Demangle.hpp>
#include <libmaus2/exception/LibMausException.hpp>

namespace libmaus2
{
	namespace parallel
	{
		namespace threadpool
		{
			struct ThreadWorkPackageData
			{
				typedef ThreadWorkPackageData this_type;
				typedef std::unique_ptr<this_type> unique_ptr_type;
				typedef std::shared_ptr<this_type> shared_ptr_type;

				virtual ~ThreadWorkPackageData() {}
				virtual std::string getPackageTypeName() const
				{
					return libmaus2::demangle::Demangle::demangleName(typeid(*this).name());
				}
				virtual std::string toString() const
				{
					return getPackageTypeName();
				}

				std::type_index getTypeIndex() const
				{
					return std::type_index(typeid(*this));
				}
			};
		}
	}
}

namespace libmaus2
{
	namespace parallel
	{
		namespace threadpool
		{
			struct ThreadWorkPackage
			{
				typedef ThreadWorkPackage this_type;
				typedef std::unique_ptr<this_type> unique_ptr_type;
				typedef std::shared_ptr<this_type> shared_ptr_type;

				uint64_t prio;
				std::pair<uint64_t,uint64_t> id;
				ThreadWorkPackageData::shared_ptr_type data;

				ThreadWorkPackage() : prio(0), id(0,0), data()
				{}

				bool operator<(ThreadWorkPackage const & T) const
				{
					// process higher priority first
					if ( prio != T.prio )
						return prio > T.prio;
					// process lower id first
					else
						return id < T.id;
				}

				ThreadWorkPackage & operator=(ThreadWorkPackage const & O) = delete;

				virtual ~ThreadWorkPackage() {}
			};
		}
	}
}

namespace libmaus2
{
	namespace parallel
	{
		namespace threadpool
		{
			struct ThreadWorkPackageQueueComparator
			{
				bool operator()(ThreadWorkPackage::shared_ptr_type const & A, ThreadWorkPackage::shared_ptr_type const & B) const
				{
					assert ( A );
					assert ( B );
					return *B < *A;
				}
			};
		}
	}
}

namespace libmaus2
{
	namespace parallel
	{
		namespace threadpool
		{
			template<typename _type>
			struct ThreadPoolStack
			{
				typedef _type type;
				typedef ThreadPoolStack<_type> this_type;
				typedef std::unique_ptr<this_type> unique_ptr_type;
				typedef std::shared_ptr<this_type> shared_ptr_type;

				std::mutex lock;
				std::stack<type> S;

				ThreadPoolStack() : lock(), S() {}

				bool pop(type & T)
				{
					std::lock_guard<std::mutex> slock(lock);
					if ( S.empty() )
						return false;

					T = S.top();
					S.pop();

					return true;
				}

				void push(type T)
				{
					std::lock_guard<std::mutex> slock(lock);
					S.push(T);
				}
			};
		}
	}
}

namespace libmaus2
{
	namespace parallel
	{
		namespace threadpool
		{
			template<typename _type>
			struct ThreadPoolQueue
			{
				typedef _type type;
				typedef ThreadPoolQueue<_type> this_type;
				typedef std::unique_ptr<this_type> unique_ptr_type;
				typedef std::shared_ptr<this_type> shared_ptr_type;

				std::mutex lock;
				std::queue<type> S;

				ThreadPoolQueue() : lock(), S() {}

				bool pop(type & T)
				{
					std::lock_guard<std::mutex> slock(lock);
					if ( S.empty() )
						return false;

					T = S.front();
					S.pop();

					return true;
				}

				void push(type T)
				{
					std::lock_guard<std::mutex> slock(lock);
					S.push(T);
				}
			};
		}
	}
}

namespace libmaus2
{
	namespace parallel
	{
		namespace threadpool
		{
			template<typename _type, typename _comparator>
			struct ThreadPoolPriorityQueue
			{
				typedef _type type;
				typedef _comparator comparator;
				typedef ThreadPoolPriorityQueue<_type,_comparator> this_type;
				typedef std::unique_ptr<this_type> unique_ptr_type;
				typedef std::shared_ptr<this_type> shared_ptr_type;

				std::mutex lock;
				std::priority_queue<type,std::vector<type>,comparator> S;

				ThreadPoolPriorityQueue() : lock(), S() {}

				bool pop(type & T)
				{
					std::lock_guard<std::mutex> slock(lock);
					if ( S.empty() )
						return false;

					T = S.top();
					S.pop();

					return true;
				}

				void push(type T)
				{
					std::lock_guard<std::mutex> slock(lock);
					S.push(T);
				}
			};
		}
	}
}

namespace libmaus2
{
	namespace parallel
	{
		namespace threadpool
		{
			template<typename value_type>
			struct LockedMapDefaultAllocator
			{
				std::shared_ptr<value_type> allocate() const
				{
					std::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;

				std::shared_ptr< std::atomic<value_type> > allocate() const
				{
					std::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;

				std::mutex lock;
				std::map<key_type,std::shared_ptr<value_type> > M;
				std::shared_ptr<allocator_type> allocator;

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

				std::shared_ptr<value_type> get(key_type const & key)
				{
					std::lock_guard<std::mutex> slock(lock);
					typename std::map<key_type,std::shared_ptr<value_type> >::iterator it = M.find(key);

					if ( it == M.end() )
					{
						std::shared_ptr<value_type> ptr(allocator->allocate());
						M[key] = ptr;
						it = M.find(key);
					}

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

					return it->second;
				}
			};
		}
	}
}

namespace libmaus2
{
	namespace parallel
	{
		namespace threadpool
		{
			struct ThreadPool;

			struct ThreadPoolDispatcher
			{
				typedef ThreadPoolDispatcher this_type;
				typedef std::unique_ptr<this_type> unique_ptr_type;
				typedef std::shared_ptr<this_type> shared_ptr_type;

				virtual ~ThreadPoolDispatcher() {}
				virtual void dispatch(ThreadWorkPackage & package, ThreadPool * pool) = 0;
			};
		}
	}
}

namespace libmaus2
{
	namespace parallel
	{
		namespace threadpool
		{
			struct ThreadPool
			{
				std::mutex dataMapLock;
				std::unordered_map<std::type_index,ThreadPoolStack<ThreadWorkPackageData::shared_ptr_type>::shared_ptr_type> dataMap;
				ThreadPoolStack<ThreadWorkPackage::shared_ptr_type> packageStack;

				std::mutex Qlock;
				std::condition_variable Qcond;
				std::priority_queue<ThreadWorkPackage::shared_ptr_type,std::vector<ThreadWorkPackage::shared_ptr_type>, ThreadWorkPackageQueueComparator> Q;

				std::atomic<int> terminatedFlag;
				std::atomic<int> panicFlag;

				std::unordered_map<std::type_index,ThreadPoolDispatcher::shared_ptr_type> dispatcherMap;

				ThreadPoolStack<libmaus2::exception::LibMausException::shared_ptr_type> panicStack;

				std::vector< std::shared_ptr<std::thread> > Vthread;

				ThreadPool(uint64_t const threads)
				:
					dataMapLock(), dataMap(), packageStack(),
					Qlock(), Qcond(), Q(),
					terminatedFlag(0),
					panicFlag(0),
					Vthread(threads)
				{
					try
					{
						for ( uint64_t i = 0; i < threads; ++i )
						{
							std::shared_ptr<std::thread> ptr(new std::thread(staticDispatch,this));
							Vthread[i] = ptr;
						}
					}
					catch(...)
					{
						panicFlag.store(1);
						throw;
					}
				}

				~ThreadPool()
				{
					terminatedFlag.store(1);
					join(false);
				}

				void terminate()
				{
					terminatedFlag = 1;
				}

				template<typename data_type>
				void registerDispatcher(ThreadPoolDispatcher::shared_ptr_type ptr)
				{
					dispatcherMap[std::type_index(typeid(data_type))] = ptr;
				}

				static void staticDispatch(ThreadPool * ptr)
				{
					ptr->dispatch();
				}

				void dispatch()
				{
					try
					{
						while ( true )
						{
							if ( panicFlag.load() )
							{
								break;
							}

							ThreadWorkPackage::shared_ptr_type pack;
							bool const ok = dequeue(pack);

							if ( ok )
							{
								assert ( pack );
								assert ( pack->data );

								std::unordered_map<std::type_index,ThreadPoolDispatcher::shared_ptr_type>::iterator it =
									dispatcherMap.find(pack->data->getTypeIndex());

								if ( it == dispatcherMap.end() )
								{
									std::string const packageTypeName = pack->data->getPackageTypeName();
									putPackage(pack);

									libmaus2::exception::LibMausException lme;
									lme.getStream() << "[E] ThreadPool::dispatch: no dispatcher found for " << packageTypeName << std::endl;
									lme.finish();
									throw lme;
								}
								else
								{
									it->second->dispatch(*pack,this);
									putPackage(pack);
								}

							}
							else if ( terminatedFlag.load() )
								break;
						}
					}
					catch(libmaus2::exception::LibMausException const & ex)
					{
						panicFlag.store(1);

						try
						{
							libmaus2::exception::LibMausException::unique_ptr_type uptr(ex.uclone());
							libmaus2::exception::LibMausException::shared_ptr_type ptr(uptr.release());
							panicStack.push(ptr);
						}
						catch(...)
						{

						}
					}
					catch(std::exception const & ex)
					{
						panicFlag.store(1);

						try
						{
							libmaus2::exception::LibMausException::shared_ptr_type ptr(new libmaus2::exception::LibMausException(ex.what()));
							panicStack.push(ptr);
						}
						catch(...)
						{

						}
					}
				}

				void join(bool const dothrow = true)
				{
					for ( uint64_t i = 0; i < Vthread.size(); ++i )
						if ( Vthread[i] && Vthread[i]->joinable() )
						{
							Vthread[i]->join();
							Vthread[i] = std::shared_ptr<std::thread>();
						}

					if ( dothrow && isInPanicMode() )
					{
						libmaus2::exception::LibMausException::shared_ptr_type exptr;
						bool const ok = panicStack.pop(exptr);

						if ( ok )
							throw *exptr;
						else
						{
							throw libmaus2::exception::LibMausException("[E] ThreadPool::join: panic mode without exception on stack");
						}
					}
				}

				bool isInPanicMode()
				{
					return panicFlag.load();
				}

				void enqueue(ThreadWorkPackage::shared_ptr_type ptr)
				{
					std::lock_guard<std::mutex> slock(Qlock);
					Q.push(ptr);
					Qcond.notify_one();
				}

				bool dequeue(ThreadWorkPackage::shared_ptr_type & ptr)
				{
					std::unique_lock<std::mutex> slock(Qlock);

					if ( Q.empty() )
						Qcond.wait_for(slock,std::chrono::seconds(1));

					if ( Q.empty() )
						return false;

					ptr = Q.top();
					Q.pop();
					return true;
				}

				ThreadWorkPackage::shared_ptr_type getRawPackage()
				{
					ThreadWorkPackage::shared_ptr_type ptr;
					bool const ok = packageStack.pop(ptr);

					if ( ok )
					{
						assert ( ptr );
						return ptr;
					}

					ThreadWorkPackage::shared_ptr_type tptr(new ThreadWorkPackage);
					return tptr;
				}

				void putRawPackage(ThreadWorkPackage::shared_ptr_type ptr)
				{
					packageStack.push(ptr);
				}

				ThreadPoolStack<ThreadWorkPackageData::shared_ptr_type>::shared_ptr_type getDataMap(std::type_index const & index)
				{
					std::lock_guard<std::mutex> slock(dataMapLock);

					std::unordered_map<std::type_index,ThreadPoolStack<ThreadWorkPackageData::shared_ptr_type>::shared_ptr_type>::iterator it = dataMap.find(index);

					if ( it == dataMap.end() )
					{
						ThreadPoolStack<ThreadWorkPackageData::shared_ptr_type>::shared_ptr_type nobj(new ThreadPoolStack<ThreadWorkPackageData::shared_ptr_type>);
						dataMap[index] = nobj;
						it = dataMap.find(index);
					}

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

					return it->second;
				}

				template<typename data_type>
				ThreadWorkPackageData::shared_ptr_type getData()
				{
					std::type_index const index(typeid(data_type));
					ThreadPoolStack<ThreadWorkPackageData::shared_ptr_type>::shared_ptr_type dmap(getDataMap(index));

					ThreadWorkPackageData::shared_ptr_type ptr;
					bool const ok = dmap->pop(ptr);

					if ( ok )
					{
						assert ( ptr );
						return ptr;
					}

					ThreadWorkPackageData::shared_ptr_type tptr(new data_type);
					return tptr;
				}

				void putData(ThreadWorkPackageData::shared_ptr_type ptr)
				{
					std::type_index const index = ptr->getTypeIndex();
					ThreadPoolStack<ThreadWorkPackageData::shared_ptr_type>::shared_ptr_type dmap(getDataMap(index));
					dmap->push(ptr);
				}

				template<typename data_type>
				ThreadWorkPackage::shared_ptr_type getPackage(uint64_t const rprio = 0, std::pair<uint64_t,uint64_t> const & id = std::pair<uint64_t,uint64_t>(0,0))
				{
					ThreadWorkPackageData::shared_ptr_type data(getData<data_type>());
					assert ( data );
					ThreadWorkPackage::shared_ptr_type pack(getRawPackage());
					assert ( pack );
					pack->prio = rprio;
					pack->id = id;
					pack->data = data;
					return pack;
				}

				void putPackage(ThreadWorkPackage::shared_ptr_type ptr)
				{
					ThreadWorkPackageData::shared_ptr_type data = ptr->data;
					ptr->data = ThreadWorkPackageData::shared_ptr_type();
					assert ( data );
					putData(data);
					assert ( ptr );
					putRawPackage(ptr);
				}
			};
		}
	}
}
#endif
