/*
    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_BGZF_DECOMPRESSTHREAD_HPP)
#define LIBMAUS2_PARALLEL_THREADPOOL_BGZF_DECOMPRESSTHREAD_HPP

#include <libmaus2/parallel/threadpool/bgzf/BGZFInputThread.hpp>

#include <zlib.h>

namespace libmaus2
{
	namespace parallel
	{
		namespace threadpool
		{
			namespace bgzf
			{
				struct BgzfBlockInfoEnqueDecompressHandler;

				struct BgzfBlockInfoProcessor
				{
					virtual ~BgzfBlockInfoProcessor() {}
				};

				struct BgzfDataDecompressedInterface
				{
					virtual ~BgzfDataDecompressedInterface() {}
					virtual void dataDecompressed(
						libmaus2::parallel::threadpool::input::ThreadPoolInputInfo * inputinfo,
						BgzfBlockInfo::shared_ptr_type blockinfo,
						uint64_t const blockid,
						uint64_t const subid,
						uint64_t const absid,
						std::shared_ptr<char[]> ddata,
						uint64_t const n,
						BgzfBlockInfoEnqueDecompressHandler * decompressHandler,
						BgzfBlockInfoProcessor * processor
					) = 0;
				};

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

					std::atomic<libmaus2::parallel::threadpool::input::ThreadPoolInputInfo *> inputinfo;
					std::atomic<BgzfDataDecompressedInterface *> callback;
					BgzfBlockInfo::shared_ptr_type blockinfo;
					std::shared_ptr<char[]> data;
					std::atomic<BgzfBlockInfoEnqueDecompressHandler *> decompHandler;

					BgzfDecompressPackageData()
					: libmaus2::parallel::threadpool::ThreadWorkPackageData(),
					  inputinfo(nullptr), callback(nullptr),
					  blockinfo(), data(),
					  decompHandler(nullptr)
					{

					}
				};

				struct BgzfDecompressContext
				{
					typedef BgzfDecompressContext this_type;
					typedef std::shared_ptr<this_type> shared_ptr_type;

					std::shared_ptr<z_stream> stream;

					BgzfDecompressContext()
					: stream(new z_stream)
					{
						std::memset(stream.get(),0,sizeof(z_stream));
						stream->zalloc = nullptr;
						stream->zfree = nullptr;
						stream->opaque = nullptr;
						stream->avail_in = 0;
						stream->next_in = nullptr;
						int const r = inflateInit2(stream.get(),-15);
						if ( r != Z_OK )
						{
							throw libmaus2::exception::LibMausException("[E] BgzfDecompressContext: inflateInit2 failed");
						}
					}

					~BgzfDecompressContext()
					{
						inflateEnd(stream.get());
					}

					void decompress(
						char * in,
						size_t n_in,
						char * out,
						size_t n_out
					)
					{
						assert ( stream );
						if ( inflateReset(stream.get()) != Z_OK )
							throw libmaus2::exception::LibMausException("[E] BgzfDecompressContext::decompress: inflateReset failed");

						stream->avail_in = n_in;
						stream->next_in = reinterpret_cast<Bytef*>(in);
						stream->avail_out = n_out;
						stream->next_out = reinterpret_cast<Bytef*>(out);

						int const r = inflate(stream.get(),Z_FINISH);

						bool const r_ok = (r == Z_STREAM_END);
						bool const out_ok = (stream->avail_out == 0);
						bool const in_ok = (stream->avail_in == 0);
						bool const ok = r_ok && out_ok && in_ok;

						if ( !ok )
						{
							throw ::libmaus2::exception::LibMausException("[E] BgzfDecompressContext::decompress: inflate failed");
						}
					}
				};

				struct BgzfDecompressedBlock
				{
					uint64_t blockid;
					uint64_t subid;
					uint64_t absid;
					std::shared_ptr<char[]> ddata;
					uint64_t n;
					BgzfBlockInfo::shared_ptr_type blockinfo;

					BgzfDecompressedBlock()
					: blockid(0), subid(0), absid(0), ddata(), n(0)
					{}
					BgzfDecompressedBlock(
						uint64_t const rblockid,
						uint64_t const rsubid,
						uint64_t const rabsid,
						std::shared_ptr<char[]> rddata,
						uint64_t const rn,
						BgzfBlockInfo::shared_ptr_type rblockinfo
					) : blockid(rblockid), subid(rsubid), absid(rabsid), ddata(rddata), n(rn), blockinfo(rblockinfo)
					{

					}
				};

				struct BgzfDecompressedBlockComparator
				{
					bool operator()(BgzfDecompressedBlock const & A, BgzfDecompressedBlock const & B) const
					{
						return A.absid > B.absid;
					}
				};


				struct BgzfBlockInfoEnqueDecompressHandler : BgzfBlockInfoHandlerInterface
				{
					struct DecompressWaitingObject
					{
						typedef std::map<uint64_t,std::shared_ptr<std::atomic<uint64_t> > > map_type;
						typedef std::shared_ptr<map_type> map_ptr_type;

						map_ptr_type Mwaiting;
						map_ptr_type Mdone;
						std::shared_ptr<std::mutex> lock;

						DecompressWaitingObject()
						: Mwaiting(new map_type), Mdone(new map_type), lock(new std::mutex)
						{
						}
						DecompressWaitingObject(DecompressWaitingObject const & O)
						: Mwaiting(O.Mwaiting), Mdone(O.Mdone), lock(O.lock)
						{}

						DecompressWaitingObject & operator=(DecompressWaitingObject const & O)
						{
							Mwaiting = O.Mwaiting;
							Mdone = O.Mdone;
							lock = O.lock;
							return *this;
						}

						void incrementWaiting(uint64_t const blockid, uint64_t const num)
						{
							std::lock_guard<std::mutex> slock(*lock);

							map_type::iterator it = Mwaiting->find(blockid);

							if ( it == Mwaiting->end() )
							{
								(*Mwaiting)[blockid] = std::shared_ptr<std::atomic<uint64_t> >(
									new std::atomic<uint64_t>(0)
								);
								it = Mwaiting->find(blockid);
							}

							assert ( it != Mwaiting->end() );

							(*(it->second)) += num;
						}

						bool incrementDone(uint64_t const blockid)
						{
							std::lock_guard<std::mutex> slock(*lock);

							map_type::iterator it = Mdone->find(blockid);

							if ( it == Mdone->end() )
							{
								(*Mdone)[blockid] = std::shared_ptr<std::atomic<uint64_t> >(
									new std::atomic<uint64_t>(0)
								);
								it = Mdone->find(blockid);
							}

							assert ( it != Mdone->end() );

							uint64_t const ldone = ++(*(it->second));

							map_type::iterator w_it = Mwaiting->find(blockid);

							assert ( w_it != Mwaiting->end() );

							return ldone == w_it->second->load();
						}
					};

					typedef libmaus2::parallel::threadpool::ThreadPoolPriorityQueue<
						BgzfDecompressedBlock,
						BgzfDecompressedBlockComparator
					> outQueueType;

					struct OutputQueueInfo
					{
						outQueueType::shared_ptr_type Q;
						std::shared_ptr<std::atomic<uint64_t> > Qnext;
						std::shared_ptr<std::mutex > Qlock;

						OutputQueueInfo()
						{
						}

						OutputQueueInfo(
							outQueueType::shared_ptr_type rQ,
							std::shared_ptr<std::atomic<uint64_t> > rQnext,
							std::shared_ptr<std::mutex > rQlock
						) : Q(rQ), Qnext(rQnext), Qlock(rQlock) {}

						void push(BgzfDecompressedBlock block)
						{
							std::lock_guard<std::mutex> slock(*Qlock);
							Q->push(block);
						}

						void next()
						{
							std::lock_guard<std::mutex> slock(*Qlock);
							++(*Qnext);
						}

						bool pop(
							BgzfDecompressedBlock & block
						)
						{
							std::lock_guard<std::mutex> slock(*Qlock);

							if ( Q->pop(block) )
							{
								if ( block.absid == Qnext->load() )
								{
									return true;
								}
								else
								{
									Q->push(block);
									return false;
								}
							}

							return false;
						}
					};

					typedef libmaus2::parallel::threadpool::ThreadPoolPriorityQueue<
						BgzfBlockInfo::shared_ptr_type,
						BgzfBlockInfoHeapComparator
					> inputQueueType;

					struct InputQueueInfo
					{
						inputQueueType::shared_ptr_type Q;
						std::shared_ptr<std::atomic<uint64_t> > Qnext;
						std::shared_ptr<std::mutex > Qlock;

						InputQueueInfo()
						{
						}

						InputQueueInfo(
							inputQueueType::shared_ptr_type rQ,
							std::shared_ptr<std::atomic<uint64_t> > rQnext,
							std::shared_ptr<std::mutex > rQlock
						) : Q(rQ), Qnext(rQnext), Qlock(rQlock) {}

						void push(BgzfBlockInfo::shared_ptr_type block)
						{
							std::lock_guard<std::mutex> slock(*Qlock);
							Q->push(block);
						}

						void next()
						{
							std::lock_guard<std::mutex> slock(*Qlock);
							++(*Qnext);
						}

						bool pop(
							BgzfBlockInfo::shared_ptr_type & block
						)
						{
							std::lock_guard<std::mutex> slock(*Qlock);

							if ( Q->pop(block) )
							{
								if ( block->absid.load() == Qnext->load() )
								{
									return true;
								}
								else
								{
									Q->push(block);
									return false;
								}
							}

							return false;
						}
					};

					struct InputQueueElement
					{
						inputQueueType::shared_ptr_type Q;
						std::shared_ptr<std::atomic<uint64_t> > next;
						std::shared_ptr<std::mutex> lock;

						InputQueueElement()
						: Q(new inputQueueType), next(new std::atomic<uint64_t>(0)), lock(new std::mutex)
						{

						}
					};

					struct OutputQueueElement
					{
						outQueueType::shared_ptr_type Q;
						std::shared_ptr<std::atomic<uint64_t> > next;
						std::shared_ptr<std::mutex> lock;

						OutputQueueElement()
						: Q(new outQueueType), next(new std::atomic<uint64_t>(0)), lock(new std::mutex)
						{

						}
					};

					ThreadPool & TP;
					libmaus2::parallel::threadpool::bgzf::ThreadPoolBGZFReadControl & TPBRC;

					typedef libmaus2::parallel::threadpool::ThreadPoolStack< std::shared_ptr<char[]> > block_stack_type;
					typedef block_stack_type::shared_ptr_type block_stack_ptr_type;

					BgzfDataDecompressedInterface * decompressedHandler;
					libmaus2::parallel::threadpool::ThreadPoolStack<BgzfDecompressContext::shared_ptr_type> contextStack;
					uint64_t numstreams;
					std::atomic<uint64_t> numStreamsReadFinished;

					libmaus2::parallel::threadpool::LockedMap<
						uint64_t,
						std::atomic<uint64_t>,
						libmaus2::parallel::threadpool::AtomicAllocator<uint64_t>
					> Mabsblock;
					libmaus2::parallel::threadpool::LockedMap<
						uint64_t,
						std::atomic<uint64_t>,
						libmaus2::parallel::threadpool::AtomicAllocator<uint64_t>
					> Mabsblockfinished;

					std::atomic<uint64_t> numStreamsDecompressFinished;

					libmaus2::parallel::threadpool::LockedMap<uint64_t,InputQueueElement> MinputQueue;
					libmaus2::parallel::threadpool::LockedMap<uint64_t,OutputQueueElement> MoutputQueue;

					BgzfBlockInfoProcessor * processor;

					libmaus2::parallel::threadpool::LockedMap<uint64_t,DecompressWaitingObject> MoutputressWaiting;

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

						struct BGZFBlockAllocator
						{
							std::shared_ptr<char[]> allocate() const
							{
								std::shared_ptr<char[]> tblock(new char[BgzfHeaderDecoderBase::getMaxBlockSize()]);
								return tblock;
							}
						};

						uint64_t const numblocks;
						BGZFBlockAllocator suballocator;

						BlockStackAllocator(uint64_t const rnumblocks)
						: numblocks(rnumblocks)
						{

						}

						block_stack_ptr_type allocate() const
						{
							block_stack_ptr_type ptr(new block_stack_type);

							for ( uint64_t i = 0; i < numblocks; ++i )
							{
								std::shared_ptr<char[]> block(suballocator.allocate());
								ptr->push(block);
							}

							return ptr;
						}
					};

					libmaus2::parallel::threadpool::LockedMap<uint64_t, std::mutex> MinputQueuePairLocks;
					BlockStackAllocator::shared_ptr_type blockStackAllocator;
					libmaus2::parallel::threadpool::LockedMap<uint64_t, block_stack_type, BlockStackAllocator> MblockStack;

					std::shared_ptr<std::mutex> getInputQueuePairLock(uint64_t const streamid)
					{
						return MinputQueuePairLocks.get(streamid);
					}

					block_stack_ptr_type getBlockStack(uint64_t const streamid)
					{
						return MblockStack.get(streamid);
					}

					DecompressWaitingObject & getDecompressWaiting(uint64_t const streamid)
					{
						return *MoutputressWaiting.get(streamid);
					}

					OutputQueueInfo getOutputQueue(uint64_t const streamid)
					{
						std::shared_ptr<OutputQueueElement> ptr = MoutputQueue.get(streamid);
						return OutputQueueInfo(ptr->Q,ptr->next,ptr->lock);
					}

					InputQueueInfo getInputQueue(uint64_t const streamid)
					{
						std::shared_ptr<InputQueueElement> ptr = MinputQueue.get(streamid);
						return InputQueueInfo(ptr->Q,ptr->next,ptr->lock);
					}

					std::atomic<uint64_t> & getAbsBlock(uint64_t const id)
					{
						return *(Mabsblock.get(id));
					}

					std::atomic<uint64_t> & getAbsBlockFinished(uint64_t const id)
					{
						return *(Mabsblockfinished.get(id));
					}

					void setDecompressedHandler(
						BgzfDataDecompressedInterface * rdecompressedHandler
					)
					{
						decompressedHandler = rdecompressedHandler;
					}

					void setProcessor(BgzfBlockInfoProcessor * rprocessor)
					{
						processor = rprocessor;
					}

					BgzfDecompressContext::shared_ptr_type getContext()
					{
						BgzfDecompressContext::shared_ptr_type ptr;

						if ( contextStack.pop(ptr) )
							return ptr;

						BgzfDecompressContext::shared_ptr_type tptr(new BgzfDecompressContext);

						return tptr;
					}

					void putContext(BgzfDecompressContext::shared_ptr_type ptr)
					{
						contextStack.push(ptr);
					}

					BgzfBlockInfoEnqueDecompressHandler(
						ThreadPool & rTP,
						libmaus2::parallel::threadpool::bgzf::ThreadPoolBGZFReadControl & rTPBRC,
						uint64_t const rnumstreams,
						uint64_t const numdecblocks = 64
					)
					: TP(rTP), TPBRC(rTPBRC), decompressedHandler(nullptr), numstreams(rnumstreams), numStreamsReadFinished(0), numStreamsDecompressFinished(0), processor(nullptr),
					  blockStackAllocator(new BlockStackAllocator(numdecblocks)), MblockStack(blockStackAllocator)
					{
					}

					virtual void decompressData(
						libmaus2::parallel::threadpool::ThreadWorkPackageData::shared_ptr_type tdata
					)
					{
						BgzfDecompressPackageData * data = dynamic_cast<BgzfDecompressPackageData *>(tdata.get());
						BgzfDecompressContext::shared_ptr_type context(getContext());

						assert ( context );

						char * from = data->blockinfo->from;
						char * to   = data->blockinfo->to;

						uint32_t const crc = BgzfHeaderDecoderBase::get32(to-2*sizeof(uint32_t),0);
						uint32_t const blocklen = BgzfHeaderDecoderBase::get32(data->blockinfo->to-1*sizeof(uint32_t),0);

						std::shared_ptr<char[]> ddata = data->data;

						assert ( ddata );

						uint32_t const headerlen = BgzfHeaderDecoder::getHeaderLength(from,0);

						context->decompress(
							from + headerlen,
							(to-from)-headerlen-2*sizeof(uint32_t),
							ddata.get(),
							blocklen
						);

						uint32_t lcrc = crc32(0,0,0);
						lcrc = crc32(lcrc, reinterpret_cast<Bytef const *>(ddata.get()), blocklen);

						if ( lcrc != crc )
						{
							throw libmaus2::exception::LibMausException("[E] BgzfBlockInfoEnqueDecompressHandler::decompressData: crc mismatch");
						}

						putContext(context);

						libmaus2::parallel::threadpool::input::ThreadPoolInputInfo * inputinfo = data->inputinfo;
						BgzfBlockInfo::shared_ptr_type blockinfo = data->blockinfo;
						BgzfBlockSummary::shared_ptr_type summary = blockinfo->owner;
						libmaus2::parallel::threadpool::input::ThreadPoolInputBlock::shared_ptr_type iblock = summary->ptr;

						uint64_t const blockid = iblock->blockid;
						uint64_t const subid = blockinfo->subid.load();
						uint64_t const absid = blockinfo->absid.load();

						#if 0
						bool const blockfinished =
						#endif
							// TPBRC.BITC.putInfo(inputinfo->streamid,blockinfo);

						data->callback.load()->dataDecompressed(
							inputinfo,
							blockinfo,
							blockid,
							subid,
							absid,
							ddata,
							blocklen,
							this,
							processor
						);
					}

					virtual void registerBgzfBlock(uint64_t const streamid, uint64_t const blockid, uint64_t const numblocks)
					{
						DecompressWaitingObject & decwa = getDecompressWaiting(streamid);
						// std::cerr << "registerBgzfBlock(streamid=" << streamid << ",blockid=" << blockid << ",numblocks=" << numblocks << ")" << std::endl;
						decwa.incrementWaiting(blockid,numblocks);
					}

					bool getDecompressionBlockPair(
						uint64_t const streamid,
						BgzfBlockInfo::shared_ptr_type & bptr,
						std::shared_ptr<char[]> & dptr
					)
					{
						InputQueueInfo IQI = getInputQueue(streamid);
						block_stack_ptr_type bstack = getBlockStack(streamid);

						std::shared_ptr<std::mutex> plock = getInputQueuePairLock(streamid);
						std::lock_guard<std::mutex> slock(*plock);

						bool const bptrok = IQI.pop(bptr);

						if ( ! bptrok )
							return false;

						bool const dptrok = bstack->pop(dptr);

						if ( ! dptrok )
						{
							IQI.push(bptr);
							return false;
						}

						return true;
					}


					void checkInputQueue(uint64_t const streamid)
					{
						InputQueueInfo IQI = getInputQueue(streamid);
						block_stack_ptr_type bstack = getBlockStack(streamid);

						BgzfBlockInfo::shared_ptr_type ptr;
						std::shared_ptr<char[]> dptr;

						while ( getDecompressionBlockPair(streamid,ptr,dptr) )
						{
							libmaus2::parallel::threadpool::ThreadWorkPackage::shared_ptr_type pack(TP.getPackage<BgzfDecompressPackageData>());
							BgzfDecompressPackageData * data = dynamic_cast<BgzfDecompressPackageData *>(pack->data.get());

							data->inputinfo = ptr->owner->inputinfo;
							data->callback = decompressedHandler;
							data->blockinfo = ptr;
							data->data = dptr;
							data->decompHandler = this;

							++getAbsBlock(ptr->owner->inputinfo->streamid);

							TP.enqueue(pack);

							IQI.next();
						}
					}

					void returnBlock(uint64_t const streamid, std::shared_ptr<char[]> block)
					{
						block_stack_ptr_type bstack = getBlockStack(streamid);
						bstack->push(block);
						checkInputQueue(streamid);
					}

					virtual void handleBgzfBlockInfo(BgzfBlockInfo::shared_ptr_type ptr)
					{
						uint64_t const streamid = ptr->owner->inputinfo->streamid;
						uint64_t const absid = ptr->absid.load();

						InputQueueInfo IQI = getInputQueue(streamid);
						IQI.push(ptr);

						checkInputQueue(streamid);
					}
				};

			}
		}
	}
}

namespace libmaus2
{
	namespace parallel
	{
		namespace threadpool
		{
			namespace bgzf
			{
				struct BgzfDecompressPackageDataDispatcher : public libmaus2::parallel::threadpool::ThreadPoolDispatcher
				{
					virtual void dispatch(libmaus2::parallel::threadpool::ThreadWorkPackage & package, libmaus2::parallel::threadpool::ThreadPool * /* pool */)
					{
						BgzfDecompressPackageData * data = dynamic_cast<BgzfDecompressPackageData *>(package.data.get());
						BgzfBlockInfoEnqueDecompressHandler * handler = data->decompHandler;
						handler->decompressData(package.data);
					}
				};
			}
		}
	}
}

namespace libmaus2
{
	namespace parallel
	{
		namespace threadpool
		{
			namespace bgzf
			{
				struct ThreadPoolBGZFDecompressControl
				{
					libmaus2::parallel::threadpool::ThreadPool & TP;
					libmaus2::parallel::threadpool::bgzf::ThreadPoolBGZFReadControl TPBRC;
					libmaus2::parallel::threadpool::bgzf::BgzfBlockInfoEnqueDecompressHandler BBITRH;

					ThreadPoolBGZFDecompressControl(
						libmaus2::parallel::threadpool::ThreadPool & rTP,
						std::vector<std::string> const & Vfn
					) : TP(rTP), TPBRC(TP,Vfn), BBITRH(TP,TPBRC,Vfn.size())
					{
						TPBRC.setBlockHandler(&BBITRH);
						TP.registerDispatcher<libmaus2::parallel::threadpool::bgzf::BgzfDecompressPackageData>(
							std::shared_ptr<libmaus2::parallel::threadpool::ThreadPoolDispatcher>(
								new libmaus2::parallel::threadpool::bgzf::BgzfDecompressPackageDataDispatcher
							)
						);
					}

					ThreadPoolBGZFDecompressControl(
						libmaus2::parallel::threadpool::ThreadPool & rTP,
						std::vector< std::shared_ptr<std::istream> > rVstr
					) : TP(rTP), TPBRC(TP,rVstr), BBITRH(TP,TPBRC,rVstr.size())
					{
						TPBRC.setBlockHandler(&BBITRH);
						TP.registerDispatcher<libmaus2::parallel::threadpool::bgzf::BgzfDecompressPackageData>(
							std::shared_ptr<libmaus2::parallel::threadpool::ThreadPoolDispatcher>(
								new libmaus2::parallel::threadpool::bgzf::BgzfDecompressPackageDataDispatcher
							)
						);
					}

					void setDecompressedHandler(
						BgzfDataDecompressedInterface * rdecompressedHandler
					)
					{
						BBITRH.setDecompressedHandler(rdecompressedHandler);
					}

					void start()
					{
						TPBRC.start();
					}

					void setProcessor(BgzfBlockInfoProcessor * rprocessor)
					{
						BBITRH.setProcessor(rprocessor);
					}

				};
			}
		}
	}
}
#endif
