/**
    bambam
    Copyright (C) 2019 German Tischler

    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/>.
**/
#include <libmaus2/lz/BgzfInflateBase.hpp>
#include <libmaus2/bambam/BamHeader.hpp>
#include <libmaus2/bambam/SamDecoder.hpp>
#include <libmaus2/util/ArgInfo.hpp>
#include <libmaus2/util/LineBuffer.hpp>
#include <libmaus2/util/Concat.hpp>
#include <libmaus2/bambam/BamBlockWriterBaseFactory.hpp>
#include <libmaus2/bambam/BamNumericalIndexDecoder.hpp>
#include <libmaus2/lz/BgzfDeflate.hpp>
#include <libmaus2/util/GetFileSize.hpp>
#include <libmaus2/util/PrefixSums.hpp>
#include <libmaus2/aio/TempFileArray.hpp>

#include <chrono>
#include <thread>

#include <biobambam2/BamBamConfig.hpp>
#include <biobambam2/Licensing.hpp>
#include <biobambam2/ReadHeader.hpp>

int getDefaultLevel()
{
	return -1;
}

int getDefaultAppend()
{
	return 0;
}

void printHelpMessage(libmaus2::util::ArgInfo const & /* arginfo */)
{
	std::cerr << ::biobambam2::Licensing::license();
	std::cerr << std::endl;
	std::cerr << "Key=Value pairs:" << std::endl;
	std::cerr << std::endl;

	std::vector< std::pair<std::string,std::string> > V;

	V.push_back ( std::pair<std::string,std::string> ( "level=<["+::biobambam2::Licensing::formatNumber(getDefaultLevel())+"]>", libmaus2::bambam::BamBlockWriterBaseFactory::getBamOutputLevelHelpText() ) );
	V.push_back ( std::pair<std::string,std::string> ( "append=<["+::biobambam2::Licensing::formatNumber(getDefaultAppend())+"]>", "append given text to header instead of replacing header by it (default: 0/replace)" ) );

	::biobambam2::Licensing::printMap(std::cerr,V);

	std::cerr << std::endl;
}

typedef std::vector<libmaus2::bambam::BamRawDecoderBase::RawInterval> offset_vector;

offset_vector getOffsets(
	std::istream & istr,
	bool const needheader,
	uint64_t const maxget = std::numeric_limits<uint64_t>::max()
)
{
	libmaus2::bambam::BamRawDecoderNoThrowEOF BRD(istr,needheader);

	std::pair<
		std::pair<uint8_t const *,uint64_t>,
		libmaus2::bambam::BamRawDecoderBase::RawInterval
	> P;

	std::vector<libmaus2::bambam::BamRawDecoderBase::RawInterval> V;

	while ( V.size() < maxget && (P=BRD.getPos()).first.first )
		V.push_back(P.second);

	return V;
}

offset_vector getOffsets(
	std::pair<uint64_t,std::string> const & compdata,
	bool const needheader,
	uint64_t const maxget = std::numeric_limits<uint64_t>::max()
)
{
	std::string const indata = compdata.second + libmaus2::lz::BgzfDeflateBase::getEOFBlock();
	std::istringstream istr(indata);
	return getOffsets(istr,needheader,maxget);
}

offset_vector getOffsets(
	std::string const & fn,
	uint64_t const maxget = std::numeric_limits<uint64_t>::max()
)
{
	libmaus2::aio::InputStreamInstance ISI(fn);
	return getOffsets(ISI,true /* need header */,maxget);
}

typedef std::pair<offset_vector,offset_vector> offset_map;

offset_map getOffsetMap(
	std::pair<uint64_t,std::string> const & compdata,
	bool const needheader,
	std::string const & bamfn
)
{
	offset_vector OA = getOffsets(compdata,needheader);
	offset_vector OB = getOffsets(bamfn,OA.size());

	assert ( OB.size() == OA.size() );

	return offset_map(OB,OA);
}

void updateMap(
	offset_map const & O,
	std::string const & indexfn,
	uint64_t const coutp,
	int64_t const shift
)
{
	std::string const tmpfn = indexfn + ".rewrite.tmp";
	libmaus2::util::TempFileRemovalContainer::addTempFile(tmpfn);
	libmaus2::aio::InputStreamInstance::unique_ptr_type ISI(new libmaus2::aio::InputStreamInstance(indexfn));
	libmaus2::aio::OutputStreamInstance::unique_ptr_type OSI(new libmaus2::aio::OutputStreamInstance(tmpfn));

	while ( ISI->peek() != std::istream::traits_type::eof() )
	{
		uint64_t const k = libmaus2::util::NumberSerialisation::deserialiseNumber(*ISI);
		uint64_t const blockpointer = libmaus2::util::NumberSerialisation::deserialiseNumber(*ISI);
		uint64_t const blockoffset = libmaus2::util::NumberSerialisation::deserialiseNumber(*ISI);

		if ( k < O.first.size() )
		{
			assert ( blockpointer == O.first[k].start.first  );
			assert ( blockoffset  == O.first[k].start.second );

			libmaus2::util::NumberSerialisation::serialiseNumber(*OSI,k);
			libmaus2::util::NumberSerialisation::serialiseNumber(*OSI,coutp + O.second[k].start.first);
			libmaus2::util::NumberSerialisation::serialiseNumber(*OSI,O.second[k].start.second);
		}
		else
		{
			libmaus2::util::NumberSerialisation::serialiseNumber(*OSI,k);
			libmaus2::util::NumberSerialisation::serialiseNumber(*OSI,static_cast<uint64_t>(static_cast<int64_t>(blockpointer)+shift));
			libmaus2::util::NumberSerialisation::serialiseNumber(*OSI,blockoffset);
		}
	}

	ISI.reset();
	OSI->flush();
	OSI.reset();

	libmaus2::aio::OutputStreamFactoryContainer::rename(tmpfn,indexfn);
}

/**
 * concatenate a set of BAM files sharing the same header (sequence dictionary) information
 * without recompressing (most of) the data blocks
 *
 * The program does NOT check the input data for actually sharing compatible headers
 *
 * The main purpose of this program is to perform a gather type operation in a scatter/gather approach
 *
 * The first argument is the file name of a text file containing input BAM file names
 *
 * An optional second argument can be used to provide a SAM file header which will be used
 * as a replacement in the output. If not given then the header from the first input BAM file
 * will be used. If the append flag is set (append=1), then the given SAM header will be appended
 * to the first BAM file header instead of replacing it
 **/
int bamfastcat(libmaus2::util::ArgInfo const & arginfo)
{
	uint64_t const numrest = arginfo.getNumRestArgs();
	int const level = libmaus2::bambam::BamBlockWriterBaseFactory::checkCompressionLevel(arginfo.getValue("level",getDefaultLevel()));
	unsigned int numthreads = arginfo.getValue<unsigned int>("threads",1);
	std::string const numerical = arginfo.getUnparsedValue("numerical",std::string());

	std::string const tmpfilenamebase = arginfo.getValue<std::string>("tmpfile",arginfo.getDefaultTmpFileName());
	std::string const tmpfilenameout = tmpfilenamebase + "_bamfastcat";

	if ( numrest < 1 )
	{
		libmaus2::exception::LibMausException lme;
		lme.getStream() << "[E] usage: " << arginfo.progname << " <in.list> [<header.sam>]" << std::endl;
		lme.finish();
		throw lme;
	}
	std::string const listfn = arginfo.getUnparsedRestArg(0);

	std::string headertext;
	if ( numrest > 1 )
	{
		std::string const headerfn = arginfo.getUnparsedRestArg(1);
		headertext = libmaus2::util::GetFileSize::readFileAsStream(headerfn);
	}

	bool const append = (headertext.size() == 0) || (arginfo.getValue("append",getDefaultAppend()) != 0);

	std::vector<std::string> Vinfn;
	{
		libmaus2::aio::InputStreamInstance ISI(listfn);
		std::string line;
		while ( std::getline(ISI,line) )
			if ( line.size() )
				Vinfn.push_back(line);
	}

	biobambam2::ReHeaderUpdate RH(headertext,append);

	libmaus2::aio::OutputStreamInstance::unique_ptr_type pOut;
	bool const haveO = arginfo.hasArg("O");

	if ( haveO )
	{
		libmaus2::aio::OutputStreamInstance::unique_ptr_type tOut(new libmaus2::aio::OutputStreamInstance(arginfo.getUnparsedValue("O",std::string())));
		pOut = std::move(tOut);
	}

	std::ostream & outstr = pOut ? (*pOut) : std::cout;

	if ( ! Vinfn.size() )
	{
		libmaus2::lz::BgzfOutputStream out(outstr,level);
		libmaus2::bambam::BamHeader const header(headertext);
		header.serialise(out);
		out.addEOFBlock();
		out.flush();

		pOut.reset();
	}
	else
	{
		std::vector< std::atomic<int64_t>  > AVA(Vinfn.size());
		std::vector< std::atomic<uint64_t> > AVM(Vinfn.size());
		libmaus2::aio::TempFileArray::unique_ptr_type pTFA;
		uint64_t totalalignments = 0;
		uint64_t totalblocks = 0;
		uint64_t numblocksize = 0;

		if ( numerical.size() )
		{
			std::atomic<int> VAall(1);

			#if defined(_OPENMP)
			#pragma omp parallel for schedule(dynamic,1) num_threads(numthreads)
			#endif
			for ( uint64_t i = 0; i < Vinfn.size(); ++i )
			{
				std::string const & bamfn = Vinfn[i];
				std::string const indexfn = libmaus2::bambam::BamNumericalIndexBase::getIndexName(bamfn);

				if ( libmaus2::util::GetFileSize::getFileSize(indexfn) )
				{
					libmaus2::bambam::BamNumericalIndexDecoder indexdec(indexfn);
					AVA[i] = indexdec.getAlignmentCount();
					AVM[i] = indexdec.getBlockSize();
				}
				else
				{
					AVA[i] = std::numeric_limits<int64_t>::min();
					AVM[i] = 0;
					VAall = 0;
				}
			}


			// if all input files have a numerical index
			if ( VAall )
			{
				std::vector<int64_t> VA(AVA.begin(),AVA.end());
				VA.push_back(0);
				libmaus2::util::PrefixSums::prefixSums(VA.begin(),VA.end());
				bool modcons = true;

				for ( uint64_t i = 1; modcons && i < Vinfn.size(); ++i )
					if ( AVM[i] != AVM[0] )
						modcons = false;

				// if all numerical indexes feature the same block size
				if ( modcons )
				{
					uint64_t const blocksize = AVM[0];

					std::string const tmppref = tmpfilenameout + "_fileparts";
					libmaus2::aio::TempFileArray::unique_ptr_type tTFA(
						new libmaus2::aio::TempFileArray(tmppref,Vinfn.size())
					);
					pTFA = std::move(tTFA);

					for ( uint64_t i = 0; i < Vinfn.size(); ++i )
					{
						uint64_t const shift = VA[i];
						std::string const & bamfn = Vinfn[i];
						std::string const indexfn = libmaus2::bambam::BamNumericalIndexBase::getIndexName(bamfn);
						libmaus2::bambam::BamNumericalIndexDecoder indexdec(indexfn);

						// number of alignments in file
						uint64_t const n = indexdec.getAlignmentCount();
						// number of alignments per thread
						uint64_t const nt = (n + numthreads - 1) / numthreads;
						// number of work packages
						uint64_t const packs = nt ? ((n+nt-1)/nt) : 0;

						totalalignments += n;

						std::string const tmpprefsub = tmpfilenameout + "_subparts";
						libmaus2::aio::TempFileArray STFA(tmpprefsub,packs);

						#if defined(_OPENMP)
						#pragma omp parallel for schedule(dynamic,1) num_threads(numthreads)
						#endif
						for ( uint64_t j = 0; j < packs; ++j )
						{
							std::ostream & out = STFA[j];
							uint64_t const from = std::min(j * nt,n);
							uint64_t const to = std::min(from+nt,n);
							uint64_t const size = to-from;
							assert ( size );
							std::pair<
								std::pair<uint8_t const *,uint64_t>,
								libmaus2::bambam::BamRawDecoderBase::RawInterval
								> P;

							libmaus2::bambam::BamRawDecoder::unique_ptr_type pdec(indexdec.getRawDecoderAt(bamfn,from));

							for ( uint64_t k = from; k < to; ++k )
							{
								P = pdec->getPos();
								assert ( P.first.first );

								if ( ((k + shift) % blocksize) == 0 )
								{
									libmaus2::util::NumberSerialisation::serialiseNumber(out,k);
									libmaus2::util::NumberSerialisation::serialiseNumber(out,P.second.start.first);
									libmaus2::util::NumberSerialisation::serialiseNumber(out,P.second.start.second);
								}
							}

							out.flush();
						}

						STFA.concat((*pTFA)[i]);
					}

					pTFA->flush();
					pTFA->close();

					totalblocks = (totalalignments + blocksize - 1)/blocksize;
					numblocksize = blocksize;
				}
			}
		}

		libmaus2::lz::BgzfInflateBase BIB;
		uint64_t coutp = 0;

		{
			std::string const & bamfn = Vinfn[0];
			libmaus2::aio::InputStreamInstance ISI(bamfn);
			std::pair<uint64_t,std::string> const compdata = readHeader(ISI,RH,level,false /* drop */);

			if ( pTFA )
			{
				offset_map const OM = getOffsetMap(compdata,true /* need header */,bamfn);
				updateMap(
					OM,
					pTFA->getName(0),
					coutp,
					static_cast<int64_t>(compdata.second.size())-
					static_cast<int64_t>(compdata.first)+
					static_cast<int64_t>(coutp)
				);
			}

			outstr.write(compdata.second.c_str(),compdata.second.size());
			coutp += compdata.second.size();

			while ( ISI.peek() != std::istream::traits_type::eof() )
			{
				libmaus2::lz::BgzfInflateBase::BaseBlockInfo BBI = BIB.readBlock(ISI);
				if ( BBI.uncompdatasize )
					coutp += BIB.writeBlock(outstr,BBI);
			}
		}

		for ( uint64_t i = 1 ; i < Vinfn.size(); ++i )
		{
			std::string const & bamfn = Vinfn[i];
			libmaus2::aio::InputStreamInstance ISI(bamfn);
			std::pair<uint64_t,std::string> const compdata = readHeader(ISI,RH,level,true /* drop */);

			if ( pTFA )
			{
				offset_map const OM = getOffsetMap(compdata,false /* need header */,bamfn);
				updateMap(
					OM,
					pTFA->getName(i),
					coutp,
					static_cast<int64_t>(compdata.second.size())-
					static_cast<int64_t>(compdata.first)+
					static_cast<int64_t>(coutp)
				);
			}

			outstr.write(compdata.second.c_str(),compdata.second.size());
			coutp += compdata.second.size();

			while ( ISI.peek() != std::istream::traits_type::eof() )
			{
				libmaus2::lz::BgzfInflateBase::BaseBlockInfo BBI = BIB.readBlock(ISI);
				if ( BBI.uncompdatasize )
					coutp += BIB.writeBlock(outstr,BBI);
			}
		}

		std::string const eofblock = libmaus2::lz::BgzfDeflateBase::getEOFBlock();
		outstr.write(eofblock.c_str(),eofblock.size());
		coutp += eofblock.size();

		outstr.flush();

		pOut.reset();

		if ( pTFA && numerical.size() )
		{
			// sleep for 5 seconds so index does not appear to be older than BAM file
			// this only makes sense if we do not write to the standard output channel
			if ( haveO )
				std::this_thread::sleep_for (std::chrono::seconds(5));

			libmaus2::aio::OutputStreamInstance::unique_ptr_type pOSI(
				new libmaus2::aio::OutputStreamInstance(numerical)
			);

			libmaus2::util::NumberSerialisation::serialiseNumber(*pOSI,totalalignments);
			libmaus2::util::NumberSerialisation::serialiseNumber(*pOSI,numblocksize);

			uint64_t blockcnt = 0;

			for ( uint64_t i = 0; i < pTFA->size(); ++i )
			{
				libmaus2::aio::InputStreamInstance::unique_ptr_type Preader(pTFA->getReader(i));

				while ( Preader->peek() != std::istream::traits_type::eof() )
				{
					libmaus2::util::NumberSerialisation::deserialiseNumber(*Preader);
					uint64_t const blockpointer = libmaus2::util::NumberSerialisation::deserialiseNumber(*Preader);
					uint64_t const blockoffset = libmaus2::util::NumberSerialisation::deserialiseNumber(*Preader);

					libmaus2::util::NumberSerialisation::serialiseNumber(*pOSI,blockpointer);
					libmaus2::util::NumberSerialisation::serialiseNumber(*pOSI,blockoffset);

					blockcnt += 1;
				}
			}

			assert ( blockcnt == totalblocks );

			pOSI->flush();
			pOSI.reset();
		}
	}

	return EXIT_SUCCESS;
}

int main(int argc, char * argv[])
{
	try
	{
		libmaus2::util::ArgInfo const arginfo(argc,argv);

		for ( uint64_t i = 0; i < arginfo.restargs.size(); ++i )
			if (
				arginfo.restargs[i] == "-v"
				||
				arginfo.restargs[i] == "--version"
			)
			{
				std::cerr << ::biobambam2::Licensing::license();
				return EXIT_SUCCESS;
			}
			else if (
				arginfo.restargs[i] == "-h"
				||
				arginfo.restargs[i] == "--help"
			)
			{
				printHelpMessage(arginfo);
				return EXIT_SUCCESS;
			}

		return bamfastcat(arginfo);
	}
	catch(std::exception const & ex)
	{
		std::cerr << ex.what() << std::endl;
		return EXIT_FAILURE;
	}
}
