/*
    libmaus2
    Copyright (C) 2021 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_BAM_THREADPOOLBAMPARSERCONTROL_HPP)
#define LIBMAUS2_PARALLEL_THREADPOOL_BAM_THREADPOOLBAMPARSERCONTROL_HPP

#include <libmaus2/util/GetCObject.hpp>
#include <libmaus2/parallel/threadpool/bam/BamParseContext.hpp>
#include <libmaus2/parallel/threadpool/bam/ThreadPoolBamParseHandler.hpp>
#include <libmaus2/parallel/threadpool/bam/BamParsePackageDataDispatcher.hpp>
#include <libmaus2/parallel/threadpool/bgzf/ThreadPoolBGZFDecompressControl.hpp>

namespace libmaus2
{
	namespace parallel
	{
		namespace threadpool
		{
			namespace bam
			{
				struct ThreadPoolBamParseControl :
					public BamParsePackageDataHandler,
					public libmaus2::parallel::threadpool::bgzf::BgzfDataDecompressedInterface
				{
					struct BgzfQueue
					{
						libmaus2::parallel::AtomicPtrHeap<libmaus2::parallel::threadpool::bgzf::BgzfDecompressedBlock> Q;
						std::atomic<std::size_t> Qnext;
						std::mutex Qlock;

						BgzfQueue()
						: Q(), Qnext(0), Qlock()
						{
						}

						void push(libmaus2::util::shared_ptr<libmaus2::parallel::threadpool::bgzf::BgzfDecompressedBlock> ptr)
						{
							std::lock_guard<decltype(Qlock)> slock(Qlock);
							Q.push(ptr);
						}

						bool pop(libmaus2::util::shared_ptr<libmaus2::parallel::threadpool::bgzf::BgzfDecompressedBlock> & ptr)
						{
							std::lock_guard<decltype(Qlock)> slock(Qlock);

							bool const ok = Q.pop(ptr);

							if ( !ok )
								return false;

							if ( ptr->absid != Qnext.load() )
							{
								Q.push(ptr);
								return false;
							}

							return true;
						}

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

					libmaus2::parallel::threadpool::ThreadPool & TP;
					libmaus2::parallel::threadpool::bgzf::ThreadPoolBGZFDecompressControl TPBDC;

					std::mutex Mcontextlock;
					libmaus2::avl::AtomicAVLPtrValueMap<std::size_t,BamParseContext> Mcontext;

					ThreadPoolBamParseHandler * const outputHandler;

					std::size_t const numstreams;

					std::atomic<std::size_t> parseFinished;

					std::shared_ptr< std::atomic<bool>[] > Afinished;

					libmaus2::avl::AtomicAVLPtrValueMap<std::size_t,BgzfQueue> Qmap;
					std::mutex Qmaplock;

					BgzfQueue & getQueue(std::size_t const streamid)
					{
						std::lock_guard<decltype(Qmaplock)> slock(Qmaplock);

						auto it = Qmap.find(streamid);

						if ( it == Qmap.end() )
						{
							libmaus2::util::shared_ptr<BgzfQueue> ptr(new BgzfQueue);
							Qmap.insert(streamid,ptr);
							it = Qmap.find(streamid);
						}

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

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

					static std::shared_ptr< std::atomic<bool>[] > constructAFinished(std::size_t const numstreams)
					{
						std::shared_ptr< std::atomic<bool>[] > Afinished(new std::atomic<bool>[numstreams]);
						for ( std::size_t i = 0; i < numstreams; ++i )
							Afinished[i].store(false);
						return Afinished;
					}

					void registerPackageDispatcher()
					{
						TP.registerDispatcher<libmaus2::parallel::threadpool::bam::BamParsePackageData>(
							libmaus2::util::shared_ptr<libmaus2::parallel::threadpool::ThreadPoolDispatcher>(
								new libmaus2::parallel::threadpool::bam::BamParsePackageDataDispatcher
							)
						);
					}

					ThreadPoolBamParseControl(
						libmaus2::parallel::threadpool::ThreadPool & rTP,
						std::vector<std::string> const & Vfn,
						ThreadPoolBamParseHandler * routputHandler,
						std::size_t const inputblocks,
						std::size_t const inputblocksize,
						std::size_t const numdecompressedblocks
					) : TP(rTP), TPBDC(TP,Vfn,this,inputblocks,inputblocksize,numdecompressedblocks), outputHandler(routputHandler), numstreams(TPBDC.TPBRC.BITC.ITC.numstreams), parseFinished(0),
					    Afinished(constructAFinished(numstreams))
					{
						registerPackageDispatcher();
					}

					ThreadPoolBamParseControl(
						libmaus2::parallel::threadpool::ThreadPool & rTP,
						std::vector< std::shared_ptr<std::istream> > rVstr,
						ThreadPoolBamParseHandler * routputHandler,
						std::size_t const inputblocks,
						std::size_t const inputblocksize,
						std::size_t const numdecompressedblocks
					) : TP(rTP), TPBDC(TP,rVstr,this,inputblocks,inputblocksize,numdecompressedblocks), outputHandler(routputHandler), numstreams(TPBDC.TPBRC.BITC.ITC.numstreams), parseFinished(0),
					    Afinished(constructAFinished(numstreams))
					{
						registerPackageDispatcher();
					}

					bool notifyEOF(std::size_t const streamid)
					{
						Afinished[streamid].store(true);
						std::size_t const lparseFinished = ++parseFinished;
						return lparseFinished == numstreams;
					}

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

					BamParseContext & getContext(std::size_t const streamid)
					{
						std::lock_guard<decltype(Mcontextlock)> slock(Mcontextlock);
						auto it = Mcontext.find(streamid);

						if ( it == Mcontext.end() )
						{
							Mcontext.insert(streamid,libmaus2::util::shared_ptr<BamParseContext>(new BamParseContext));
							it = Mcontext.find(streamid);
						}

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


					/**
					 * parse BAM
					 **/
					void handle(BamParsePackageData & data)
					{
						uint64_t const streamid = data.ddata->streamid.load();

						BgzfQueue & OQI = getQueue(streamid);

						BamParseContext & context = getContext(streamid);

						OQI.push(
							libmaus2::util::shared_ptr<libmaus2::parallel::threadpool::bgzf::BgzfDecompressedBlock>
							(
								new libmaus2::parallel::threadpool::bgzf::BgzfDecompressedBlock(
									streamid,
									data.blockid,
									data.subid,
									data.absid,
									data.ddata,
									data.n,
									data.blockinfo
								)
							)
						);

						libmaus2::util::shared_ptr<libmaus2::parallel::threadpool::bgzf::BgzfDecompressedBlock> block;
						std::atomic<std::size_t> d(0);
						// get next block in line
						while ( OQI.pop(block) )
						{
							bool const blockeof = block->blockinfo.load()->eof;

							if ( context.putback_f.load() )
								block->ddata.load()->prepend(
									context.putback.load().get(),
									context.putback.load().get() + context.putback_f.load()
								);

							context.putback_f.store(0);

							block->ddata.load()->P.reset();

							char const * ca = block->ddata.load()->B.load().get() + block->ddata.load()->o.load();
							char const * ce = ca + block->ddata.load()->f.load();

							if ( blockeof )
								std::cerr << "found bam eof on " << streamid << std::endl;

							if ( ! context.haveHeader )
							{
								libmaus2::util::GetCObject<uint8_t const *> G(reinterpret_cast<uint8_t const *>(ca));
								std::ptrdiff_t const av = ce - ca;
								auto const P = context.headerParser.parseHeader(G,av);

								if ( P.first )
								{
									context.haveHeader.store(true);
									// std::cerr << "got bam header on " << streamid << std::endl;
								}

								ca += P.second;
							}

							if ( context.haveHeader )
							{
								std::size_t blocksize;
								while (
									ce-ca >= static_cast<std::ptrdiff_t>(sizeof(uint32_t))
									&&
									ce-ca >= static_cast<std::ptrdiff_t>(sizeof(uint32_t)
										+
										#if defined(LIBMAUS2_HAVE_i386)
										(blocksize=*reinterpret_cast<uint32_t const *>(ca)))
										#else
										(blocksize=libmaus2::parallel::threadpool::bgzf::BgzfHeaderDecoderBase::get32(ca,0)))
										#endif
								)
								{
									block->ddata.load()->P.push(ca);

									ca += sizeof(uint32_t) + blocksize;

									// std::cerr << "blocksize=" << blocksize << std::endl;

									context.numalgn++;
								}

								std::size_t const n_putback = ce-ca;

								block->ddata.load()->f -= n_putback;

								context.addPutback(ca,ce);
							}

							// check whether data is truncated
							if ( context.putback_f.load() && blockeof )
							{
								libmaus2::exception::LibMausException lme;
								lme.getStream() << "[E] ThreadPoolBamParseControl::handle(): unexpected EOF in decompressed stream" << std::endl;
								lme.finish();
								throw lme;
							}

							// enque block for processing
							context.putOutputQueue(block->ddata);

							assert ( block->ddata.load()->decompressHandler == data.decompressHandler );

							// return compressed block data
							data.decompressHandler.load()->TPBRC.BITC.putInfo(streamid,block->blockinfo);

							// call handler for parsed data
							outputHandler->handle(this,streamid);

							if ( context.numalgn.load() >> 20 != context.numalgnprnt.load() >> 20 )
							{
								std::cerr << streamid << "\t" << context.numalgn.load() << std::endl;

								context.numalgnprnt.store(context.numalgn.load());
							}

							// ready for next block
							OQI.next();

							++d;
						}
					}

					/**
					 * handle decompressed data and enque parsing package
					 **/
					virtual void dataDecompressed(
						// libmaus2::parallel::threadpool::input::ThreadPoolInputInfo * inputinfo,
						libmaus2::util::shared_ptr<libmaus2::parallel::threadpool::bgzf::BgzfBlockInfo> blockinfo,
						uint64_t const blockid,
						uint64_t const subid,
						uint64_t const absid,
						libmaus2::util::shared_ptr<libmaus2::parallel::threadpool::bgzf::GenericBlock> ddata,
						uint64_t const n
					)
					{
						auto decompressHandler = ddata->decompressHandler.load();

						libmaus2::util::shared_ptr<libmaus2::parallel::threadpool::ThreadWorkPackage> pack(TP.getPackage<BamParsePackageData>());
						BamParsePackageData * data = dynamic_cast<BamParsePackageData *>(pack->data.load().get());

						// data->inputinfo = inputinfo;
						data->blockinfo = blockinfo;
						data->blockid = blockid;
						data->subid = subid;
						data->absid = absid;
						data->ddata = ddata;
						data->n = n;
						data->decompressHandler = decompressHandler;
						data->handler.store(this);

						TP.enqueue(pack);
					}

				};
			}
		}
	}
}
#endif
