/*
    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
				{
					ThreadPool & TP;
					libmaus2::parallel::threadpool::bgzf::ThreadPoolBGZFReadControl & TPBRC;
					// BgzfInputThreadControl & BITC;
					libmaus2::parallel::threadpool::ThreadPoolStack< std::shared_ptr<char[]> > blockStack;
					BgzfDataDecompressedInterface * decompressedHandler;
					libmaus2::parallel::threadpool::ThreadPoolStack<BgzfDecompressContext::shared_ptr_type> contextStack;
					uint64_t numstreams;
					std::atomic<uint64_t> numStreamsReadFinished;

					std::map<uint64_t,std::shared_ptr<std::atomic<uint64_t> > > Mabsblock;
					std::mutex Mabsblocklock;
					std::map<uint64_t,std::shared_ptr<std::atomic<uint64_t> > > Mabsblockfinished;
					std::mutex Mabsblockfinishedlock;

					std::atomic<uint64_t> numStreamsDecompressFinished;

					typedef libmaus2::parallel::threadpool::ThreadPoolPriorityQueue<
						BgzfDecompressedBlock,
						BgzfDecompressedBlockComparator
					> outQueueType;
					std::mutex decompQueueLock;
					std::map<uint64_t,outQueueType::shared_ptr_type> MdecompQueue;
					std::map<uint64_t,std::shared_ptr<std::atomic<uint64_t> > > MdecompQueueNext;
					std::map<uint64_t,std::shared_ptr<std::mutex > > MdecompQueueLock;

					BgzfBlockInfoProcessor * processor;

					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();
						}
					};

					std::map<uint64_t,DecompressWaitingObject> MdecompressWaiting;
					std::mutex MdecompressWaitingLock;

					DecompressWaitingObject getDecompressWaiting(uint64_t const streamid)
					{
						std::lock_guard<std::mutex> slock(MdecompressWaitingLock);

						std::map<uint64_t,DecompressWaitingObject>::iterator it =
							MdecompressWaiting.find(streamid);

						if ( it == MdecompressWaiting.end() )
						{
							MdecompressWaiting[streamid] = DecompressWaitingObject();
							it = MdecompressWaiting.find(streamid);
						}

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

						return it->second;
					}

					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;
						}
					};

					OutputQueueInfo getOutputQueue(uint64_t const streamid)
					{
						std::lock_guard<std::mutex> slock(decompQueueLock);

						std::map<uint64_t,outQueueType::shared_ptr_type>::iterator it =
							MdecompQueue.find(streamid);

						if ( it == MdecompQueue.end() )
						{
							MdecompQueueNext[streamid] =
								std::shared_ptr<std::atomic<uint64_t> >(new std::atomic<uint64_t>(0));
							MdecompQueueLock[streamid] =
								std::shared_ptr<std::mutex >(new std::mutex);
							MdecompQueue[streamid] =
								outQueueType::shared_ptr_type(new outQueueType);

							it = MdecompQueue.find(streamid);
						}

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

						std::map<uint64_t,std::shared_ptr<std::atomic<uint64_t> > >::iterator n_it =
							MdecompQueueNext.find(streamid);
						assert ( n_it != MdecompQueueNext.end() );
						std::map<uint64_t,std::shared_ptr<std::mutex> >::iterator l_it =
							MdecompQueueLock.find(streamid);
						assert ( l_it != MdecompQueueLock.end() );

						return OutputQueueInfo(it->second,n_it->second,l_it->second);
					}

					std::atomic<uint64_t> & getAbsBlock(uint64_t const id)
					{
						std::lock_guard<std::mutex> slock(Mabsblocklock);
						std::map<uint64_t,std::shared_ptr<std::atomic<uint64_t> > >::iterator it = Mabsblock.find(id);
						assert ( it != Mabsblock.end() );
						return *(it->second);
					}

					std::atomic<uint64_t> & getAbsBlockFinished(uint64_t const id)
					{
						std::lock_guard<std::mutex> slock(Mabsblockfinishedlock);
						std::map<uint64_t,std::shared_ptr<std::atomic<uint64_t> > >::iterator it = Mabsblockfinished.find(id);
						assert ( it != Mabsblockfinished.end() );
						return *(it->second);
					}

					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);
					}

					std::shared_ptr<char[]> getBlock()
					{
						std::shared_ptr<char[]> block;

						if ( blockStack.pop(block) )
							return block;

						std::shared_ptr<char[]> tblock(new char[BgzfHeaderDecoderBase::getMaxBlockSize()]);

						return tblock;
					}

					void returnBlock(std::shared_ptr<char[]> block)
					{
						blockStack.push(block);
					}

					BgzfBlockInfoEnqueDecompressHandler(
						ThreadPool & rTP,
						libmaus2::parallel::threadpool::bgzf::ThreadPoolBGZFReadControl & rTPBRC,
						uint64_t const rnumstreams
					)
					: TP(rTP), TPBRC(rTPBRC), blockStack(), decompressedHandler(nullptr), numstreams(rnumstreams), numStreamsReadFinished(0), numStreamsDecompressFinished(0), processor(nullptr)
					{
						for ( uint64_t i = 0; i < numstreams; ++i )
						{
							Mabsblock[i] = std::shared_ptr<std::atomic<uint64_t> >(
								new std::atomic<uint64_t>(0)
							);
							Mabsblockfinished[i] = std::shared_ptr<std::atomic<uint64_t> >(
								new std::atomic<uint64_t>(0)
							);
						}
					}

					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);
					}

					virtual void handleBgzfBlockInfo(BgzfBlockInfo::shared_ptr_type ptr)
					{
						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 = getBlock();
						data->decompHandler = this;

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

						TP.enqueue(pack);
					}
				};

			}
		}
	}
}

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
