/*
    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_CRAM_THREADPOOLCRAMENQUEUECONTROL_HPP)
#define LIBMAUS2_PARALLEL_THREADPOOL_CRAM_THREADPOOLCRAMENQUEUECONTROL_HPP

#include <libmaus2/parallel/threadpool/bam/BamHeaderSeqData.hpp>
#include <libmaus2/parallel/threadpool/cram/CramEnqueuePackageDataDispatcher.hpp>
// #include <libmaus2/bambam/BamAlignmentDecoderBase.hpp>
#include <libmaus2/parallel/threadpool/bam/ThreadPoolBamParseControl.hpp>
#include <libmaus2/parallel/threadpool/cram/CramEncodingWorkPackageDispatcher.hpp>
#include <libmaus2/parallel/threadpool/bgzf/BgzfHeaderDecoderBase.hpp>

namespace libmaus2
{
	namespace parallel
	{
		namespace threadpool
		{
			namespace cram
			{
				struct ThreadPoolCramEnqueueControl : public libmaus2::parallel::threadpool::bam::ThreadPoolBamParseHandler, CramEnqueuePackageDataHandler
				{
					struct CramBlockEncodingInfo
					{
						libmaus2::util::atomic_shared_ptr<libmaus2::util::shared_ptr<libmaus2::parallel::threadpool::bgzf::GenericBlock>[] > Vblock;
						std::atomic<std::size_t> numblocks;
						libmaus2::util::atomic_shared_ptr<char const *[]> cram_block;
						libmaus2::util::atomic_shared_ptr<std::size_t []> cram_blocksize;
						libmaus2::util::atomic_shared_ptr<std::size_t []> cram_blockelements;
						std::atomic<std::size_t> inputblockid;
						std::atomic<std::size_t> numel;

						CramBlockEncodingInfo(
							libmaus2::util::shared_ptr<libmaus2::util::shared_ptr<libmaus2::parallel::threadpool::bgzf::GenericBlock>[] > rVblock,
							std::size_t const n,
							std::size_t const rinputblockid
						) : Vblock(rVblock),
						    numblocks(n),
						    cram_block(libmaus2::util::shared_ptr<char const *[]>(new char const *[n])),
						    cram_blocksize(libmaus2::util::shared_ptr<std::size_t []>(new size_t[n])),
						    cram_blockelements(libmaus2::util::shared_ptr<std::size_t []>(new size_t[n])),
						    inputblockid(rinputblockid),
						    numel(0)
						{
							for ( std::size_t i = 0; i < n; ++i )
							{
								std::size_t const lnumel = Vblock.load()[i]->P.f.load();

								if ( lnumel )
								{
									char const * first_p = Vblock.load()[i]->P.p.load()[0];
									char const * last_p  = Vblock.load()[i]->P.p.load()[lnumel-1];
									std::size_t const blocklength_last = libmaus2::parallel::threadpool::bgzf::BgzfHeaderDecoderBase::get32(last_p,0 /* offset */);
									std::size_t const data_size = (last_p-first_p)+blocklength_last+sizeof(uint32_t);

									cram_block.load()[i]         = first_p;
									cram_blocksize.load()[i]     = data_size;
									cram_blockelements.load()[i] = lnumel;
								}
								else
								{
									cram_block.load()[i] = nullptr;
									cram_blocksize.load()[i] = 0;
									cram_blockelements.load()[i] = 0;
								}

								numel += lnumel;

								#if 0
								for ( std::size_t j = 0; j < Vblock.load()[i]->P.f.load(); ++j )
								{
									char const * p = Vblock.load()[i]->P.p.load()[j];
									std::size_t const blocklength = libmaus2::parallel::threadpool::bgzf::BgzfHeaderDecoderBase::get32(p,0 /* offset */);
									std::cerr << "blocklength=" << blocklength << std::endl;
									char const * name = libmaus2::bambam::BamAlignmentDecoderBase::getReadName(reinterpret_cast<uint8_t const *>(p+sizeof(uint32_t)));
									std::cerr << name << std::endl;

									auto valid = libmaus2::bambam::BamAlignmentDecoderBase::valid(
										reinterpret_cast<uint8_t const *>(p+sizeof(uint32_t)),
										blocklength
									);

									bool const ok = valid.valid == libmaus2::bambam::libmaus2_bambam_alignment_validity_ok;
									if ( ! ok )
									{
										std::cerr << "failed " << j << " / " << Vblock.load()[i]->P.f.load() << std::endl;
										assert ( ok );
									}
								}
								#endif
							}
						}
					};

					struct OutputBlock
					{
						libmaus2::util::atomic_shared_ptr<char[]> B;
						std::atomic<std::size_t> n;
						std::atomic<std::size_t> f;
						std::atomic<int64_t> inblockid;
						std::atomic<std::size_t> outblockid;
						std::atomic<bool> blockfinal;

						OutputBlock()
						: B(), n(0), f(0), inblockid(0), outblockid(0), blockfinal(false)
						{
						}

						std::pair<char const *,char const *> get() const
						{
							char const * ca = B.load().get();
							char const * ce = ca + f.load();
							return std::make_pair(ca,ce);
						}

						void put(char const * ca, char const * ce)
						{
							std::size_t const c = ce - ca;

							if ( c > n.load() )
							{
								libmaus2::util::shared_ptr<char[]> T(new char[c]);
								B.store(T);
								n.store(c);
							}

							std::copy(ca,ce,B.load().get());
							f.store(c);
						}

						bool operator<(OutputBlock const & O) const
						{
							if ( inblockid != O.inblockid )
								return inblockid < O.inblockid;
							else
								return outblockid < O.outblockid;
						}
					};

					libmaus2::parallel::threadpool::ThreadPool & TP;
					libmaus2::parallel::threadpool::bam::ThreadPoolBamParseControl TPBPC;
					std::size_t const cramaccum;
					std::atomic<void *> cramEncoder;
					std::mutex cramEncoderLock;
					libmaus2::parallel::threadpool::bam::BamHeaderSeqData const & dict;
					std::atomic<std::size_t> nextCramInputBlock;

					libmaus2::avl::AtomicAVLPtrValueMap<std::size_t,CramBlockEncodingInfo> cramBlockInfoMap;
					std::mutex cramBlockInfoMapLock;

					libmaus2::parallel::AtomicPtrStack<OutputBlock> outputStack;
					libmaus2::parallel::AtomicPtrHeap<OutputBlock> outputHeap;
					std::mutex outputLock;
					std::pair< std::atomic<int64_t>, std::atomic<std::size_t> > outputNext;

					std::string const profile;

					std::size_t const seqsperslice;
					std::size_t const basesperslice;
					std::string const cramversion;

					bool getNextOutputBlock(libmaus2::util::shared_ptr<OutputBlock> & ptr)
					{
						std::lock_guard<decltype(outputLock)> slock(outputLock);

						bool const ok = outputHeap.pop(ptr);

						if ( ! ok )
							return false;

						if (
							ptr->inblockid != outputNext.first
							||
							ptr->outblockid != outputNext.second
						)
						{
							outputHeap.push(ptr);
							return false;
						}

						return true;
					}

					libmaus2::util::shared_ptr<OutputBlock> getOutputBlock()
					{
						libmaus2::util::shared_ptr<OutputBlock> tptr;

						if ( outputStack.pop(tptr) )
							return tptr;

						libmaus2::util::shared_ptr<OutputBlock> nptr(new OutputBlock);

						return nptr;
					}

					void putOutputBlock(libmaus2::util::shared_ptr<OutputBlock> ptr)
					{
						outputStack.push(ptr);
					}

					void * getEncoder(libmaus2::parallel::threadpool::bam::BamParseContext & context)
					{
						std::lock_guard<decltype(cramEncoderLock)> lock(cramEncoderLock);

						if ( ! cramEncoder.load() )
						{
							if ( !context.haveHeader.load() )
								return nullptr;

							auto const & parser = context.headerParser;

							char const * headertext = parser.text.begin();
							std::size_t const l_text = parser.l_text;

							std::string sheadertext(headertext,headertext + l_text);

							std::istringstream sheaderistr(sheadertext);
							libmaus2::parallel::threadpool::bam::BamHeaderSeqData ldict(sheaderistr);

							for ( std::map<std::string,std::map<std::string,std::string> >::iterator itP = ldict.M.begin(); itP != ldict.M.end(); ++itP )
							{
								std::string const & name = itP->first;
								std::map<std::string,std::string> & LM = itP->second;

								auto const it = dict.M.find(name);

								if ( it == dict.M.end() )
								{
									libmaus2::exception::LibMausException lme;
									lme.getStream() << "ThreadPoolCramEnqueueControl::getEncoder(): unable to find sequence " << name << " in sequence dictionary" << std::endl;
									lme.finish();
									throw lme;
								}

								std::map<std::string,std::string> const & OM = it->second;

								if ( LM.find("M5") == LM.end() )
								{
									auto const m5it = OM.find("M5");

									if ( m5it == OM.end() )
									{
										libmaus2::exception::LibMausException lme;
										lme.getStream() << "ThreadPoolCramEnqueueControl::getEncoder(): unable to find seq hash for sequence " << name << " in sequence dictionary" << std::endl;
										lme.finish();
										throw lme;
									}
									else
									{
										LM["M5"] = m5it->second;
									}
								}

								if ( LM.find("UR") == LM.end() )
								{
									auto const m5it = OM.find("UR");

									if ( m5it == OM.end() )
									{
										libmaus2::exception::LibMausException lme;
										lme.getStream() << "ThreadPoolCramEnqueueControl::getEncoder(): unable to find seq hash for sequence " << name << " in sequence dictionary" << std::endl;
										lme.finish();
										throw lme;
									}
									else
									{
										LM["UR"] = m5it->second;
									}
								}
							}


							{
								std::ostringstream oheaderstr;
								bool first = true;
								std::istringstream sheaderistr(sheadertext);

								std::string line;
								std::regex reg("^@SQ");

								while ( std::getline(sheaderistr,line) )
									if ( std::regex_search(line,reg) )
									{
										if ( first )
										{
											for ( auto s : ldict.NV )
											{
												auto const & M = ldict.M.find(s)->second;

												oheaderstr << "@SQ";

												if ( M.find("SN") != M.end() )
													oheaderstr << "\t" << "SN" << ":" << M.find("SN")->second;
												if ( M.find("LN") != M.end() )
													oheaderstr << "\t" << "LN" << ":" << M.find("LN")->second;
												if ( M.find("M5") != M.end() )
													oheaderstr << "\t" << "M5" << ":" << M.find("M5")->second;
												if ( M.find("UR") != M.end() )
													oheaderstr << "\t" << "UR" << ":" << M.find("UR")->second;


												#if 0
												for ( auto P : M )
													oheaderstr << "\t" << P.first << ":" << P.second;
												#endif

												oheaderstr << "\n";
											}
											first = false;
										}
									}
									else
									{
										oheaderstr << line << "\n";
									}

								sheadertext = oheaderstr.str();
							}

							// std::cerr << sheadertext;

							auto writefunc = cram_data_write_function_static;

							assert ( writefunc );


							void * encoder = libmaus2::bambam::parallel::ScramCramEncoding::io_lib_cram_allocate_encoder(this,sheadertext.c_str(),sheadertext.size(),writefunc);

							if ( ! encoder )
							{
								libmaus2::exception::LibMausException lme;
								lme.getStream() << "ThreadPoolCramEnqueueControl::getEncoder(): failed in cram_allocate_encoder" << std::endl;
								lme.finish();
								throw lme;
							}

							if ( cramversion.size() )
							{
								libmaus2::bambam::parallel::ScramCramEncoding::io_lib_cram_set_version(encoder,cramversion);
							}

							libmaus2::bambam::parallel::ScramCramEncoding::io_lib_cram_set_profile(encoder,profile);

							if ( seqsperslice != 0 )
							{
								libmaus2::bambam::parallel::ScramCramEncoding::io_lib_cram_set_seqs_per_slice(encoder,seqsperslice);
								if ( ! basesperslice )
									libmaus2::bambam::parallel::ScramCramEncoding::io_lib_cram_set_bases_per_slice(encoder,seqsperslice*500);
							}
							else if ( basesperslice != 0 )
							{
								libmaus2::bambam::parallel::ScramCramEncoding::io_lib_cram_set_bases_per_slice(encoder,basesperslice);
							}


							cramEncoder.store(encoder);
						}

						return cramEncoder.load();
					}

					void cram_data_write_function(ssize_t const inblockid, size_t const outblockid, char const *data, size_t const n, cram_data_write_block_type const blocktype)
					{
						auto outblock = getOutputBlock();

						outblock->put(data,data+n);
						outblock->inblockid.store(inblockid);
						outblock->outblockid.store(outblockid);

						switch ( blocktype )
						{
							case cram_data_write_block_type_block_final:
							case cram_data_write_block_type_file_final:
								outblock->blockfinal.store(true);
								break;
							default:
								outblock->blockfinal.store(false);
								break;
						}

						{
							std::lock_guard<decltype(outputLock)> slock(outputLock);
							outputHeap.push(outblock);
						}

						libmaus2::util::shared_ptr<OutputBlock> ptr;
						while ( getNextOutputBlock(ptr) )
						{
							auto const P = ptr->get();

							std::cout.write(P.first,P.second-P.first);

							putOutputBlock(ptr);

							{
								std::lock_guard<decltype(outputLock)> slock(outputLock);
								if ( ptr->blockfinal )
								{
									outputNext.first++;
									outputNext.second.store(0);
								}
								else
								{
									outputNext.second++;
								}
							}
						}
					}

					/* enqueue a work package */
					void cram_enque_compression_work_package_function(void * workpackage)
					{
						libmaus2::util::shared_ptr<ThreadWorkPackage> tpack(TP.getPackage<CramEncodingWorkPackageData>());
						CramEncodingWorkPackageData & pack = dynamic_cast<CramEncodingWorkPackageData &>(*(tpack->data.load()));
						pack.workpackage.store(workpackage);
						TP.enqueue(tpack);
					}

					void cram_compression_work_package_finished(size_t const inblockid, int const final)
					{
						std::lock_guard<decltype(cramBlockInfoMapLock)> slock(cramBlockInfoMapLock);
						auto it = cramBlockInfoMap.find(inblockid);
						assert ( it != cramBlockInfoMap.end() );
						auto craminfo = it->second.load();
						cramBlockInfoMap.erase(it);

						for ( std::size_t i = 0; i < craminfo->numblocks.load(); ++i )
						{
							auto block = craminfo->Vblock.load()[i];
							auto decompressHandler = block->decompressHandler.load();
							decompressHandler->returnBlock(block->streamid.load(),block);
						}

						if ( final )
						{
							TP.terminate();
						}
					}

					static void cram_data_write_function_static(void *userdata, ssize_t const inblockid, size_t const outblockid, char const *data, size_t const n, cram_data_write_block_type const blocktype)
					{
						reinterpret_cast<ThreadPoolCramEnqueueControl *>(userdata)->cram_data_write_function(inblockid,outblockid,data,n,blocktype);
					}

					static void cram_enque_compression_work_package_function_static(void * userdata, void * workpackage)
					{
						reinterpret_cast<ThreadPoolCramEnqueueControl *>(userdata)->cram_enque_compression_work_package_function(workpackage);
					}

					static void cram_compression_work_package_finished_static(void *userdata, size_t const inblockid, int const final)
					{
						reinterpret_cast<ThreadPoolCramEnqueueControl *>(userdata)->cram_compression_work_package_finished(inblockid,final);
					}

					void registerDispatcher()
					{
						TP.registerDispatcher<libmaus2::parallel::threadpool::cram::CramEnqueuePackageData>(
							libmaus2::util::shared_ptr<libmaus2::parallel::threadpool::ThreadPoolDispatcher>(
								new libmaus2::parallel::threadpool::cram::CramEnqueuePackageDataDispatcher
							)
						);

						TP.registerDispatcher<libmaus2::parallel::threadpool::cram::CramEncodingWorkPackageData>(
							libmaus2::util::shared_ptr<libmaus2::parallel::threadpool::ThreadPoolDispatcher>(
								new libmaus2::parallel::threadpool::cram::CramEncodingWorkPackageDispatcher
							)
						);
					}

					ThreadPoolCramEnqueueControl(
						libmaus2::parallel::threadpool::ThreadPool & rTP,
						std::vector< std::shared_ptr<std::istream> > rVstr,
						libmaus2::parallel::threadpool::bam::BamHeaderSeqData const & rdict,
						std::size_t const inputblocks,
						std::size_t const inputblocksize,
						std::size_t const numdecompressedblocks,
						std::size_t const cramblockaccum,
						std::string const & rprofile,
						std::size_t const rseqsperslice,
						std::size_t const rbasesperslice,
						std::string const & rcramversion
					) : TP(rTP), TPBPC(rTP,rVstr,this,inputblocks,inputblocksize,numdecompressedblocks), cramaccum(cramblockaccum), cramEncoder(nullptr), cramEncoderLock(), dict(rdict), nextCramInputBlock(0),
					    outputNext(-1,0), profile(rprofile),
					    seqsperslice(rseqsperslice),
					    basesperslice(rbasesperslice),
					    cramversion(rcramversion)
					{
						registerDispatcher();
					}

					~ThreadPoolCramEnqueueControl()
					{
						if ( cramEncoder.load() )
						{
							libmaus2::bambam::parallel::ScramCramEncoding::io_lib_cram_deallocate_encoder(cramEncoder.load());
							cramEncoder.store(nullptr);
						}
					}

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

					virtual void handle(libmaus2::parallel::threadpool::bam::ThreadPoolBamParseControl * control, std::size_t const streamid)
					{
						libmaus2::util::shared_ptr<ThreadWorkPackage> tpack(TP.getPackage<CramEnqueuePackageData>());
						CramEnqueuePackageData & pack = dynamic_cast<CramEnqueuePackageData &>(*(tpack->data.load()));

						pack.control.store(control);
						pack.streamid.store(streamid);
						pack.handler.store(this);

						TP.enqueue(tpack);
					}

					void cramEnqueue(libmaus2::parallel::threadpool::bam::ThreadPoolBamParseControl * control, std::size_t const streamid)
					{
						auto & context = control->getContext(streamid);
						libmaus2::util::shared_ptr<libmaus2::util::shared_ptr<libmaus2::parallel::threadpool::bgzf::GenericBlock>[] > Vblock;
						std::size_t n;

						while ( (n=context.getOutputQueueElements(Vblock,cramaccum)) )
						{
							bool const eof = Vblock[n-1]->eof.load();

							if ( eof && !context.haveHeader.load() )
							{
								libmaus2::exception::LibMausException lme;
								lme.getStream() << "[E] ThreadPoolCramEnqueueControl::cramEnqueue: reached end of input stream, but BAM header is not complete" << std::endl;
								lme.finish();
								throw lme;
							}

							std::size_t numel = 0;
							for ( std::size_t i = 0; i < n; ++i )
								numel += Vblock[i]->f.load();

							if ( numel || eof )
							{
								void * encoder = getEncoder(context);

								if ( ! encoder )
								{
									libmaus2::exception::LibMausException lme;
									lme.getStream() << "[E] ThreadPoolCramEnqueueControl::cramEnqueue: failed to instantiate encoder although we have data records" << std::endl;
									lme.finish();
									throw lme;
								}

								libmaus2::util::shared_ptr<CramBlockEncodingInfo> craminfo(
									new CramBlockEncodingInfo(Vblock,n,nextCramInputBlock++)
								);

								{
									std::lock_guard<decltype(cramBlockInfoMapLock)> slock(cramBlockInfoMapLock);
									cramBlockInfoMap.insert(craminfo->inputblockid.load(),craminfo);
								}

								int const rq = libmaus2::bambam::parallel::ScramCramEncoding::io_lib_cram_enque_compression_block(
									this,
									encoder,
									craminfo->inputblockid.load(),
									craminfo->cram_block.load().get(),
									craminfo->cram_blocksize.load().get(),
									craminfo->cram_blockelements.load().get(),
									n,
									eof ? 1 : 0, /* final */
									cram_enque_compression_work_package_function_static,
									cram_data_write_function_static,
									cram_compression_work_package_finished_static
								);

								if ( rq )
								{
									libmaus2::exception::LibMausException lme;
									lme.getStream() << "[E] ThreadPoolCramEnqueueControl::cramEnqueue: cram_enque_compression_block failed" << std::endl;
									lme.finish();
									throw lme;
								}
							}

							context.bumpOutNext(n);
						}
					}
				};
			}
		}
	}
}
#endif
