/*
    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/parallel/threadpool/ThreadPoolDispatcher.hpp>
#include <libmaus2/parallel/threadpool/LockedMap.hpp>
#include <libmaus2/parallel/AtomicPtrStack.hpp>
#include <libmaus2/demangle/Demangle.hpp>
#include <libmaus2/exception/LibMausException.hpp>
#include <libmaus2/avl/AtomicAVLPtrValueMap.hpp>
#include <libmaus2/avl/AtomicAVLKeyPtrSet.hpp>

namespace libmaus2
{
	namespace parallel
	{
		namespace threadpool
		{
			struct ThreadPool
			{
				uint64_t const threads;
				std::mutex dataMapLock;
				libmaus2::avl::AtomicAVLPtrValueMap<std::type_index, libmaus2::parallel::AtomicPtrStack<ThreadWorkPackageData> > dataMap;
				libmaus2::parallel::AtomicPtrStack<ThreadWorkPackage> packageStack;

				std::mutex Qlock;
				std::condition_variable Qcond;
				libmaus2::avl::AtomicAVLKeyPtrSet<ThreadWorkPackage> Q;

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

				std::mutex dispatcherMapLock;
				libmaus2::avl::AtomicAVLPtrValueMap<std::type_index,ThreadPoolDispatcher> dispatcherMap;
				libmaus2::parallel::AtomicPtrStack<libmaus2::exception::LibMausException> panicStack;
				libmaus2::parallel::AtomicPtrStack< std::thread > Vthread;

				void registerGenericDispatcher()
				{
					libmaus2::util::shared_ptr<ThreadPoolGenericDispatcher> disp(new ThreadPoolGenericDispatcher);
					registerDispatcher<ThreadPoolGenericPackageData>(disp);
				}

				ThreadPool(uint64_t const rthreads)
				:
					threads(rthreads),
					dataMapLock(), dataMap(), packageStack(),
					Qlock(), Qcond(), Q(),
					terminatedFlag(0),
					panicFlag(0),
					seq(0),
					Vthread(threads)
				{
					try
					{
						registerGenericDispatcher();

						for ( uint64_t i = 0; i < threads; ++i )
						{
							libmaus2::util::shared_ptr<std::thread> ptr(new std::thread(staticDispatch,this));
							Vthread.set(i,ptr);
						}
					}
					catch(...)
					{
						panicFlag.store(1);
						throw;
					}
				}

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

				void terminate()
				{
					terminatedFlag = 1;
				}

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

				template<typename data_type>
				void registerDispatcher(libmaus2::util::shared_ptr<ThreadPoolDispatcher> ptr)
				{
					std::lock_guard<decltype(dispatcherMapLock)> slock(dispatcherMapLock);

					std::type_index const type(typeid(data_type));

					auto const it = dispatcherMap.find(type);

					if ( it != dispatcherMap.end() )
					{
						libmaus2::exception::LibMausException lme;
						lme.getStream() << "[E] dispatcher for type is already registered" << std::endl;
						lme.finish();
						throw lme;
					}

					dispatcherMap.insert(std::type_index(typeid(data_type)),ptr);
				}

				template<typename data_type>
				void ensureDispatcherRegistered(libmaus2::util::shared_ptr<ThreadPoolDispatcher> ptr)
				{
					std::lock_guard<decltype(dispatcherMapLock)> slock(dispatcherMapLock);

					std::type_index const type(typeid(data_type));

					auto const it = dispatcherMap.find(type);

					if ( it == dispatcherMap.end() )
						dispatcherMap.insert(std::type_index(typeid(data_type)),ptr);
				}

				ThreadPoolDispatcher * getDispatcher(std::type_index const & TI)
				{
					std::lock_guard<decltype(dispatcherMapLock)> slock(dispatcherMapLock);

					auto it = dispatcherMap.find(TI);

					if ( it == dispatcherMap.end() )
						return nullptr;

					return it->second.load().get();
				}

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

							libmaus2::util::shared_ptr<ThreadWorkPackage> pack;
							bool const ok = dequeue(pack);

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

								auto * dispatcher = getDispatcher(pack->data.load()->getTypeIndex());

								if ( dispatcher )
								{
									dispatcher->dispatch(pack,this);
									putPackage(pack);
								}
								else
								{
									std::string const packageTypeName = pack->data.load()->getPackageTypeName();
									putPackage(pack);

									libmaus2::exception::LibMausException lme;
									lme.getStream() << "[E] ThreadPool::dispatch: no dispatcher found for " << packageTypeName << std::endl;
									lme.finish();
									throw lme;
								}
							}
							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::util::shared_ptr<libmaus2::exception::LibMausException> ptr(uptr.release());
							panicStack.push(ptr);
						}
						catch(...)
						{

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

						try
						{
							libmaus2::util::shared_ptr<libmaus2::exception::LibMausException> ptr(new libmaus2::exception::LibMausException(ex.what()));
							panicStack.push(ptr);
						}
						catch(...)
						{

						}
					}
				}

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

					if ( dothrow && isInPanicMode() )
					{
						libmaus2::util::shared_ptr<libmaus2::exception::LibMausException> 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(libmaus2::util::shared_ptr<ThreadWorkPackage> ptr)
				{
					std::lock_guard<std::mutex> slock(Qlock);

					Q.insert(ptr);
					Qcond.notify_one();
				}

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

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

					if ( Q.empty() )
						return false;

					auto it = Q.begin();
					assert ( it != Q.end() );

					ptr = it->load();

					Q.erase(it);

					return true;
				}

				libmaus2::util::shared_ptr<ThreadWorkPackage> getRawPackage()
				{
					libmaus2::util::shared_ptr<ThreadWorkPackage> ptr;
					bool const ok = packageStack.pop(ptr);

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

					libmaus2::util::shared_ptr<ThreadWorkPackage> tptr(new ThreadWorkPackage);
					return tptr;
				}

				void putRawPackage(libmaus2::util::shared_ptr<ThreadWorkPackage> ptr)
				{
					packageStack.push(ptr);
				}

				libmaus2::util::shared_ptr< AtomicPtrStack<ThreadWorkPackageData> > getDataMap(std::type_index const & index)
				{
					std::lock_guard<std::mutex> slock(dataMapLock);

					auto it = dataMap.find(index);

					if ( it == dataMap.end() )
					{
						libmaus2::util::shared_ptr < AtomicPtrStack<ThreadWorkPackageData> > nobj(new AtomicPtrStack<ThreadWorkPackageData>);
						dataMap.insert(index,nobj);
						it = dataMap.find(index);
					}

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

					return it->second.load();
				}

				template<typename data_type>
				libmaus2::util::shared_ptr<ThreadWorkPackageData> getData()
				{
					std::type_index const index(typeid(data_type));
					libmaus2::util::shared_ptr< AtomicPtrStack<ThreadWorkPackageData> > dmap(getDataMap(index));

					libmaus2::util::shared_ptr<ThreadWorkPackageData> ptr;
					bool const ok = dmap->pop(ptr);

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

					libmaus2::util::shared_ptr<ThreadWorkPackageData> tptr(new data_type);
					return tptr;
				}

				void putData(libmaus2::util::shared_ptr<ThreadWorkPackageData> ptr)
				{
					std::type_index const index = ptr->getTypeIndex();
					libmaus2::util::shared_ptr< AtomicPtrStack<ThreadWorkPackageData> > dmap(getDataMap(index));
					dmap->push(ptr);
				}

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

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