/*
    biobambam2
    Copyright (C) 2009-2015 German Tischler
    Copyright (C) 2011-2015 Genome Research Limited

    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/bambam/BamMultiAlignmentDecoderFactory.hpp>
#include <libmaus2/bambam/BamBlockWriterBaseFactory.hpp>
#include <libmaus2/lz/BgzfInflateBase.hpp>
#include <libmaus2/lz/BgzfDeflateZStreamBase.hpp>
#include <libmaus2/lz/BgzfDeflate.hpp>

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

static int getDefaultLevel() {return libmaus2::lz::DeflateDefaults::getDefaultLevel();}
int getDefaultVerbose() { return 0; }
std::string getDefaultInputFormat() { return "bam"; }
uint64_t getDefaultSizeThres() { return 32*1024*1024; }
std::string getDefaultPrefix() { return "split_"; }

int bamexploderef(libmaus2::util::ArgInfo const & arginfo)
{
	std::string const prefix = arginfo.getUnparsedValue("prefix",getDefaultPrefix());
	uint64_t const threads = arginfo.getValue<uint64_t>("threads",1);
	int const level = libmaus2::lz::DeflateDefaults::getDefaultLevel();

	for ( uint64_t i = 0; i < arginfo.getNumRestArgs(); ++i )
	{
		std::string const infn = arginfo.getUnparsedRestArg(i);

		libmaus2::bambam::BamIndex const index(infn);
		libmaus2::bambam::BamHeader::unique_ptr_type pheader;

		std::vector<std::string> Vref;

		// get BAM header
		{
			libmaus2::util::ArgInfo arginfoinfn = arginfo;
			arginfoinfn.replaceKey("I",infn);
			libmaus2::bambam::BamAlignmentDecoderWrapper::unique_ptr_type decwrapper(
				libmaus2::bambam::BamMultiAlignmentDecoderFactory::construct(
					arginfoinfn
				)
			);
			libmaus2::bambam::BamAlignmentDecoder & dec = decwrapper->getDecoder();

			libmaus2::bambam::BamHeader const & header = dec.getHeader();
			for ( uint64_t j = 0; j < header.getNumRef(); ++j )
			{
				std::string const refname = header.getRefIDName(j);
				Vref.push_back(refname);
			}

			libmaus2::bambam::BamHeader::unique_ptr_type theader(header.uclone());
			pheader = std::move(theader);
		}

		std::vector<std::string> Vout(Vref.size());
		std::atomic<int> gfailed(0);
		#if defined(_OPENMP)
		#pragma omp parallel for num_threads(threads) schedule(dynamic,1)
		#endif
		for ( uint64_t j = 0; j < Vref.size(); ++j )
		{
			try
			{
				libmaus2::bambam::BamRangeChromosome range(Vref[j],*pheader);
				std::vector< std::pair<uint64_t,uint64_t> > const chunks = range.getChunks(index);

				std::ostringstream fnostr;
				fnostr
					<< prefix
					<< std::setw(6) << std::setfill('0') << i << std::setw(0)
					<< "_"
					<< std::setw(6) << std::setfill('0') << j << std::setw(0)
					<< std::setw(0) << ".bam";
					;

				std::string const fn = fnostr.str();

				Vout[j] = fn;

				libmaus2::aio::InputStreamInstance ISI(infn);
				libmaus2::lz::BgzfInflateBase BIS;
				libmaus2::autoarray::AutoArray<uint8_t> B(BIS.getBgzfMaxBlockSize());
				libmaus2::autoarray::AutoArray<char> C(BIS.getBgzfMaxBlockSize());
				libmaus2::lz::BgzfDeflateOutputBufferBase outbuf(level);
				libmaus2::lz::BgzfDeflateZStreamBase outbase(level);
				libmaus2::aio::OutputStreamInstance OSI(fn);

				{
					libmaus2::lz::BgzfDeflate<std::ostream> BD(OSI,level);
					pheader->serialise(BD);
					BD.flush();
				}

				for ( uint64_t z = 0; z < chunks.size(); ++z )
				{
					if ( chunks[z].second > chunks[z].first )
					{
						uint64_t sp = (chunks[z].first>>16);
						ISI.seekg(sp);

						// do we handle only a single block?
						bool const oneblockonly = ((chunks[z].first >> 16) == (chunks[z].second >> 16));

						// handle first block
						{
							libmaus2::lz::BgzfInflateBase::BaseBlockInfo const BBI = BIS.readBlock(ISI);
							BIS.decompressBlock(reinterpret_cast<char *>(B.begin()),BBI);

							uint8_t * start = B.begin() + (chunks[z].first &0xFFFF);
							uint8_t * end   = B.begin() + (oneblockonly ? (chunks[z].second&0xFFFF) : BBI.uncompdatasize);
							libmaus2::lz::BgzfDeflateZStreamBaseFlushInfo const outinfo = outbase.flush(start,end,outbuf);

							uint64_t const n = (outinfo.blocks == 1) ? outinfo.block_a_c : (outinfo.block_a_c + outinfo.block_b_c);

							OSI.write(reinterpret_cast<char const *>(outbuf.outbuf.begin()),n);

							sp += BBI.compdatasize;
						}

						if ( ! oneblockonly )
						{
							uint64_t tocopy = (chunks[z].second>>16) - sp;

							while ( tocopy )
							{
								uint64_t const toread = std::min(tocopy,C.size());

								ISI.read(C.begin(),toread);
								assert ( ISI && ISI.gcount() == static_cast<int64_t>(toread) );
								OSI.write(C.begin(),toread);

								tocopy -= toread;
								sp += toread;
							}

							assert ( sp == (chunks[z].second >> 16) );

							if ( (chunks[z].second & 0xFFFF) != 0 )
							{
								libmaus2::lz::BgzfInflateBase::BaseBlockInfo const BBI = BIS.readBlock(ISI);
								BIS.decompressBlock(reinterpret_cast<char *>(B.begin()),BBI);

								uint8_t * start = B.begin();
								uint8_t * end   = B.begin() + (chunks[z].second&0xFFFF);
								libmaus2::lz::BgzfDeflateZStreamBaseFlushInfo const outinfo = outbase.flush(start,end,outbuf);

								uint64_t const n = (outinfo.blocks == 1) ? outinfo.block_a_c : (outinfo.block_a_c + outinfo.block_b_c);

								OSI.write(reinterpret_cast<char const *>(outbuf.outbuf.begin()),n);
							}
						}


						#if 0
						std::cerr
							<< j << " " << z << " "
							<< "from " << (chunks[z].first>>16) << "," << (chunks[z].first & 0xFFFF)
							<< " "
							<< "to " << (chunks[z].second>>16) << "," << (chunks[z].second & 0xFFFF)
							<< std::endl;
						#endif
					}

				}

				std::string const eofblock = libmaus2::lz::BgzfDeflate<std::ostream>::getEOFBlock();
				OSI.write(eofblock.c_str(),eofblock.size());

				OSI.flush();
			}
			catch(std::exception const & ex)
			{
				libmaus2::aio::StreamLock::lock_type::scope_lock_type slock(libmaus2::aio::StreamLock::cerrlock);
				std::cerr << "[E] error for sequence " << j << " (" << Vref[j] << "):" << std::endl;
				std::cerr << ex.what() << std::endl;
				gfailed = 1;
			}
		}

		if ( gfailed )
		{
			libmaus2::exception::LibMausException lme;
			lme.getStream() << "[E] bamexploderef failed in parallel loop" << std::endl;
			lme.finish();
			throw lme;
		}

		for ( uint64_t j = 0; j < Vout.size(); ++j )
		{
			std::cout << infn << "\t" << j << "\t" << Vout[j] << "\n";
		}
	}

	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"
			)
			{
				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> ( "verbose=<["+::biobambam2::Licensing::formatNumber(getDefaultVerbose())+"]>", "print progress report" ) );
				V.push_back ( std::pair<std::string,std::string> ( std::string("inputformat=<[")+getDefaultInputFormat()+"]>", std::string("input format (") + libmaus2::bambam::BamMultiAlignmentDecoderFactory::getValidInputFormats() + ")" ) );
				V.push_back ( std::pair<std::string,std::string> ( std::string("outputformat=<[")+libmaus2::bambam::BamBlockWriterBaseFactory::getDefaultOutputFormat()+"]>", std::string("output format (") + libmaus2::bambam::BamBlockWriterBaseFactory::getValidOutputFormats() + ")" ) );
				V.push_back ( std::pair<std::string,std::string> ( "I=<[stdin]>", "input filename (standard input if unset)" ) );
				V.push_back ( std::pair<std::string,std::string> ( "inputthreads=<[1]>", "input helper threads (for inputformat=bam only, default: 1)" ) );
				V.push_back ( std::pair<std::string,std::string> ( "reference=<>", "reference FastA (.fai file required, for cram i/o only)" ) );
				V.push_back ( std::pair<std::string,std::string> ( "range=<>", "coordinate range to be processed (for coordinate sorted indexed BAM input only)" ) );
				V.push_back ( std::pair<std::string,std::string> ( "outputthreads=<[1]>", "output helper threads (for outputformat=bam only, default: 1)" ) );
				V.push_back ( std::pair<std::string,std::string> ( "O=<[stdout]>", "output filename (standard output if unset)" ) );
				V.push_back ( std::pair<std::string,std::string> ( std::string("prefix=<[")+getDefaultPrefix()+"]>", "prefix of output file names" ) );

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

				std::cerr << std::endl;
				return EXIT_SUCCESS;
			}

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