/**
    bambam
    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/>.
**/
#include "config.h"

#include <iostream>
#include <queue>

#include <libmaus2/util/ArgInfo.hpp>
#include <libmaus2/util/ArgParser.hpp>
#include <libmaus2/timing/RealTimeClock.hpp>
#include <libmaus2/bambam/BamBlockWriterBaseFactory.hpp>
#include <libmaus2/bambam/BamMergeCoordinate.hpp>
#include <libmaus2/bambam/BamMergeQueryName.hpp>
#include <libmaus2/bambam/BamWriter.hpp>
#include <libmaus2/bambam/BamRawDecoder.hpp>
#include <libmaus2/bambam/BamMultiAlignmentDecoderFactory.hpp>
#include <libmaus2/lz/BgzfDeflateOutputCallbackMD5.hpp>
#include <libmaus2/bambam/BgzfDeflateOutputCallbackBamIndex.hpp>

#include <biobambam2/Licensing.hpp>

static int getDefaultVerbose() { return 1; }

::libmaus2::bambam::BamHeader::unique_ptr_type updateHeader(
	::libmaus2::util::ArgParser const & arg,
	::libmaus2::bambam::BamHeader const & header
)
{
	std::string const headertext(header.text);

	// add PG line to header
	std::string const upheadtext = ::libmaus2::bambam::ProgramHeaderLineSet::addProgramLine(
		headertext,
		"bamsalvage", // ID
		"bamsalvage", // PN
		arg.commandline, // CL
		::libmaus2::bambam::ProgramHeaderLineSet(headertext).getLastIdInChain(), // PP
		std::string(PACKAGE_VERSION) // VN
	);
	// construct new header
	::libmaus2::bambam::BamHeader::unique_ptr_type uphead(new ::libmaus2::bambam::BamHeader(upheadtext));

	return uphead;
}

/**
 * save as much data as possible from a damaged BAM file equipped with a .bai index
 **/
int bamsalvage(libmaus2::util::ArgParser const & arg)
{
	std::string const infn = arg[0];

	arg.printArgs(std::cerr);

	int const verbose = arg.getParsedArgOrDefault<int>("verbose",getDefaultVerbose());

	::libmaus2::bambam::BamHeader::shared_ptr_type sheader;

	// load BAM header
	{
		libmaus2::bambam::BamDecoder dec(infn);
		::libmaus2::bambam::BamHeader::unique_ptr_type uphead(updateHeader(arg,dec.getHeader()));
		::libmaus2::bambam::BamHeader::shared_ptr_type tptr(uphead->sclone());
		sheader = tptr;
	}

	libmaus2::bambam::BamIndex const index(infn);
	libmaus2::autoarray::AutoArray<libmaus2::bambam::BamIndexRef> const & refs = index.getRefs();


	libmaus2::bambam::BamBlockWriterBase::unique_ptr_type Pwriter(
		libmaus2::bambam::BamBlockWriterBaseFactory::construct(*sheader,arg,nullptr));

	libmaus2::aio::InputStreamInstance ISI(infn);
	uint64_t galcnt = 0;

	for ( uint64_t refid = 0; refid < refs.size(); ++refid )
	{
		if ( verbose )
			std::cerr << "[V] refid=" << refid << " sequence=" << sheader->getRefIDName(refid) << std::endl;

		auto const & ref = refs[refid];
		auto const & lin = ref.lin;
		auto const & intervals = lin.intervals;

		// extract values from linear index which are not null pointers (undefined)
		std::vector<uint64_t> Voff;
		for ( std::size_t i = 0; i < intervals.size(); ++i )
			if ( intervals[i] )
				Voff.push_back(intervals[i]);

		// check sanity of index values
		for ( std::size_t i = 1; i < Voff.size(); ++i )
			assert ( Voff[i-1] <= Voff[i] );

		for ( std::size_t i = 0; i < Voff.size(); ++i )
		{
			uint64_t const low = Voff[i];
			uint64_t const high = (i+1<Voff.size()) ? Voff[i+1] : std::numeric_limits<uint64_t>::max();

			if ( low != 0 )
			{
				libmaus2::bambam::BamRawDecoder dec(
					ISI,false,
					std::make_pair(
						low>>16,
						low&0xFFFFull
					)
				);

				std::pair<uint64_t,uint64_t> const E(
					high>>16,
					high&0xFFFFull
				);

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


					while ( (P=dec.getPos()).first.first )
					{
						if (
							P.second.start == E
							||
							(
								libmaus2::bambam::BamAlignmentDecoderBase::getRefID(P.first.first)
								!=
								static_cast<int64_t>(refid)
							)
						)
							break;

						assert (
							((P.second.start.first << 16) | (P.second.start.second))
							<
							high
						);

						Pwriter->writeBamBlock(P.first.first,P.first.second);

						uint64_t const lgalcnt = ++galcnt;

						if ( verbose && ((lgalcnt % (1024*1024)) == 0) )
							std::cerr << "[V] " << lgalcnt << std::endl;

						if ( P.second.end == E )
							break;
					}
				}
				catch(std::exception const & ex)
				{
					std::cerr
						<< "low=(" << (low>>16) << "," << (low&0xFFFF) << ") "
						<< "high=(" << (high>>16) << "," << (high&0xFFFF) << ")"
						<< std::endl;
					std::cerr << ex.what() << std::endl;
				}
			}
		}
	}

	if ( verbose )
		std::cerr << "[V] " << galcnt << std::endl;

	Pwriter.reset();

	std::cout << libmaus2::lz::BgzfDeflate<std::ostream>::getEOFBlock();
	std::cout.flush();

	return EXIT_SUCCESS;
}

int main(int argc, char * argv[])
{
	try
	{
		std::vector<libmaus2::util::ArgParser::ArgumentDefinition> Vformatcons;
		Vformatcons.push_back(libmaus2::util::ArgParser::ArgumentDefinition("v","verbose",true));

		Vformatcons.push_back(libmaus2::util::ArgParser::ArgumentDefinition("h","help",true));

		std::vector<libmaus2::util::ArgParser::ArgumentDefinition> const Vformatout = libmaus2::bambam::BamBlockWriterBaseFactory::getArgumentDefinitions();

		std::vector<libmaus2::util::ArgParser::ArgumentDefinition> Vformat = libmaus2::util::ArgParser::mergeFormat(Vformatout,Vformatcons);

		libmaus2::util::ArgParser const arg(argc,argv,Vformat);

		if ( arg.uniqueArgPresent("help") || arg.size() < 1 )
		{
			std::cerr << ::biobambam2::Licensing::license();
			std::cerr << std::endl;

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

			V.push_back ( std::pair<std::string,std::string> ( "-v/--verbose=0/1", "print progress report" ) );

			std::cerr << "usage: " << arg.progname << " [args] <in.bam>\n\n";

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

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

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