/*
    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_BGZFBLOCKINFOENQUEDECOMPRESSHANDLER_HPP)
#define LIBMAUS2_PARALLEL_THREADPOOL_BGZF_BGZFBLOCKINFOENQUEDECOMPRESSHANDLER_HPP

#include <libmaus2/parallel/threadpool/bgzf/BgzfDecompressPackageData.hpp>
#include <libmaus2/parallel/threadpool/bgzf/BgzfDecompressContext.hpp>
#include <libmaus2/parallel/threadpool/bgzf/BgzfDecompressedBlock.hpp>

namespace libmaus2
{
	namespace parallel
	{
		namespace threadpool
		{
			namespace bgzf
			{
				struct BgzfBlockInfoEnqueDecompressHandler : BgzfBlockInfoHandlerInterface
				{
					typedef libmaus2::parallel::AtomicPtrHeap<BgzfDecompressedBlock> outQueueType;

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

						OutputQueueInfo()
						{
						}

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

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

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

						bool pop(
							libmaus2::util::shared_ptr<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::AtomicPtrHeap<BgzfBlockInfo> inputQueueType;

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

						InputQueueInfo()
						{
						}

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

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

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

						bool pop(
							libmaus2::util::shared_ptr<BgzfBlockInfo> & 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
					{
						libmaus2::util::shared_ptr<inputQueueType> 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
					{
						std::shared_ptr<outQueueType> 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::AtomicPtrStack< GenericBlock > block_stack_type;
					typedef libmaus2::util::shared_ptr<block_stack_type> block_stack_ptr_type;

					BgzfDataDecompressedInterface * decompressedHandler;
					libmaus2::parallel::AtomicPtrStack<BgzfDecompressContext> 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;

					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
						{
							libmaus2::util::shared_ptr<GenericBlock> allocate() const
							{
								libmaus2::util::shared_ptr<GenericBlock> tblock(new GenericBlock(BgzfHeaderDecoderBase::getMaxBlockSize()+4096u));
								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 )
							{
								libmaus2::util::shared_ptr<GenericBlock> block(suballocator.allocate());
								ptr->push(block);
							}

							return ptr;
						}

						uint64_t getNumBlocks() const
						{
							return numblocks;
						}
					};

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

					uint64_t getDecoderBlocks() const
					{
						return blockStackAllocator->getNumBlocks();
					}

					libmaus2::util::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);
					}

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

					InputQueueInfo getInputQueue(uint64_t const streamid)
					{
						libmaus2::util::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;
					}

					libmaus2::util::shared_ptr<BgzfDecompressContext> getContext()
					{
						libmaus2::util::shared_ptr<BgzfDecompressContext> ptr;

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

						libmaus2::util::shared_ptr<BgzfDecompressContext> tptr(new BgzfDecompressContext);

						return tptr;
					}

					void putContext(libmaus2::util::shared_ptr<BgzfDecompressContext> ptr)
					{
						contextStack.push(ptr);
					}

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

					virtual void decompressData(
						libmaus2::util::shared_ptr<libmaus2::parallel::threadpool::ThreadWorkPackageData> tdata
					)
					{
						BgzfDecompressPackageData * data = dynamic_cast<BgzfDecompressPackageData *>(tdata.get());
						libmaus2::util::shared_ptr<BgzfDecompressContext> 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);

						libmaus2::util::shared_ptr<GenericBlock> ddata = data->data;

						assert ( ddata );

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

						ddata->f.store(blocklen);
						ddata->o.store(ddata->n.load() - blocklen);
						ddata->absid.store(data->blockinfo->absid);
						ddata->streamid.store(data->inputinfo.load()->streamid);
						ddata->decompressHandler.store(this);
						ddata->eof.store(data->blockinfo->eof.load());
						ddata->inputinfo.store(data->inputinfo);

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

						uint32_t const lcrc = context->crc32(ddata->B.load().get() + ddata->o.load(), blocklen);

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

						putContext(context);

						libmaus2::util::shared_ptr<BgzfBlockInfo> blockinfo = data->blockinfo;
						libmaus2::util::shared_ptr<BgzfBlockSummary> summary = blockinfo->owner;
						libmaus2::util::shared_ptr<libmaus2::parallel::threadpool::input::ThreadPoolInputBlock> iblock = summary->ptr;

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

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

					bool getDecompressionBlockPair(
						uint64_t const streamid,
						libmaus2::util::shared_ptr<BgzfBlockInfo> & bptr,
						libmaus2::util::shared_ptr<GenericBlock> & dptr
					)
					{
						InputQueueInfo IQI = getInputQueue(streamid);
						block_stack_ptr_type bstack = getBlockStack(streamid);

						libmaus2::util::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);

						libmaus2::util::shared_ptr<BgzfBlockInfo> ptr;
						libmaus2::util::shared_ptr<GenericBlock> dptr;

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

							data->inputinfo = ptr->owner.load()->inputinfo;
							data->callback = decompressedHandler;
							data->blockinfo = ptr;
							data->data = dptr;
							data->decompHandler = this;

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

							TP.enqueue(pack);

							IQI.next();
						}
					}

					void returnBlock(uint64_t const streamid, libmaus2::util::shared_ptr<GenericBlock> block)
					{
						block_stack_ptr_type bstack = getBlockStack(streamid);
						bstack->push(block);
						checkInputQueue(streamid);
					}

					virtual void handleBgzfBlockInfo(libmaus2::util::shared_ptr<BgzfBlockInfo> ptr)
					{
						uint64_t const streamid = ptr->owner.load()->inputinfo->streamid;

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

						checkInputQueue(streamid);
					}
				};

			}
		}
	}
}
#endif
