/**
    biobambam2
    Copyright (C) 2020 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 <biobambam2/BamBamConfig.hpp>
#include <biobambam2/Licensing.hpp>

#include <libmaus2/bambam/BamBlockWriterBaseFactory.hpp>
#include <libmaus2/bambam/BamMultiAlignmentDecoderFactory.hpp>
#include <libmaus2/util/stringFunctions.hpp>

#include <iomanip>
#include <iostream>

#include <libmaus2/util/ArgParser.hpp>

#include <config.h>

int bamexploderg(libmaus2::util::ArgParser const & arg)
{
	std::string const outprefix = arg[0];

	std::map<std::string,std::string> replaceM;
	::libmaus2::trie::LinearHashTrie<char,uint32_t>::shared_ptr_type LHTsnofailure;
	std::vector<std::string> replaceIn;
	std::vector<std::string> replaceOut;
	std::string const outputformat = arg("outputformat","bam");

	if ( arg.uniqueArgPresent("rgreplacetable") )
	{
		libmaus2::aio::InputStreamInstance ISI(arg["rgreplacetable"]);
		std::string line;
		while ( std::getline(ISI,line) )
		{
			std::deque<std::string> const D = libmaus2::util::stringFunctions::tokenize(line,std::string("\t"));
			if ( D.size() >= 2 && D[0].size() && D[1].size() )
			{
				if ( replaceM.find(D[0]) != replaceM.end() )
				{
					libmaus2::exception::LibMausException lme;
					lme.getStream() << "[E] read group " << D[0] << " appears multple times as source in replacement file " << arg["rgreplacetable"] << std::endl;
					lme.finish();
					throw lme;
				}
				replaceM[D[0]] = D[1];

				// std::cerr << "[R]\t" << D[0] << "\t" << D[1] << std::endl;
			}
		}

		if ( replaceM.size() )
		{
			std::vector<std::string> replaceIn;

			for ( std::map<std::string,std::string>::const_iterator it = replaceM.begin(); it != replaceM.end(); ++it )
			{
				replaceIn.push_back(it->first);
				replaceOut.push_back(it->second);
			}

			::libmaus2::trie::Trie<char> trienofailure;
			trienofailure.insertContainer(replaceIn);
			::libmaus2::trie::LinearHashTrie<char,uint32_t>::unique_ptr_type LHTnofailure(trienofailure.toLinearHashTrie<uint32_t>());
			LHTsnofailure = ::libmaus2::trie::LinearHashTrie<char,uint32_t>::shared_ptr_type(LHTnofailure.release());
		}
	}

	libmaus2::bambam::BamAlignmentDecoderWrapper::unique_ptr_type pdec(
		libmaus2::bambam::BamMultiAlignmentDecoderFactory::construct(
			arg
		)
	);

	libmaus2::bambam::BamAlignmentDecoder & decoder = pdec->getDecoder();
	libmaus2::bambam::BamHeader const & lheader = pdec->getDecoder().getHeader();
	libmaus2::bambam::BamHeader::shared_ptr_type sheader = lheader.sclone();
	libmaus2::bambam::BamAlignment & algn = decoder.getAlignment();
	uint64_t const numrg = lheader.getNumReadGroups();
	libmaus2::autoarray::AutoArray<libmaus2::bambam::BamBlockWriterBase::unique_ptr_type> Aout(numrg+1);

	libmaus2::bambam::BamAuxFilterVector RGtag;
	RGtag.set("RG");

	uint64_t lineid = 0;
	while ( decoder.readAlignment() )
	{
		char const * rg = algn.getReadGroup();
		int64_t const rgid = lheader.getReadGroupId(rg);
		assert ( rgid >= -1 );
		uint64_t const rgindex = rgid+1;

		if ( ! Aout[rgindex] )
		{
			std::string outfn;

			if ( rgid < 0 )
				outfn = outprefix + "default";
			else
			{
				std::string rgname = lheader.getReadGroupIdentifierAsString(rgid);
				std::map<std::string,std::string>::const_iterator it = replaceM.find(rgname);
				if ( it != replaceM.end() )
					rgname = it->second;

				outfn = outprefix + rgname;
			}

			outfn += std::string(".") + outputformat;

			libmaus2::util::ArgParser larg = arg;
			larg.replaceArg("O",outfn);

			// std::cerr << "[U] " << larg << std::endl;

			libmaus2::aio::OutputStreamFactoryContainer::makeDirectoryPath(outfn,0755);

			libmaus2::bambam::BamBlockWriterBase::unique_ptr_type tptr(
				libmaus2::bambam::BamBlockWriterBaseFactory::construct(lheader,larg,nullptr /* callbacks */)
			);

			std::cerr << "[F]\t" << rgid << "\t" << ((rgid >= 0)?lheader.getReadGroupIdentifierAsString(rgid):"*") << "\t" << outfn << std::endl;

			Aout.at(rgindex) = std::move(tptr);
		}

		if ( rg && LHTsnofailure )
		{
			int64_t const id = LHTsnofailure->searchCompleteNoFailureZ(rg);

			if ( id >= 0 )
			{
				std::string const replaceby = replaceOut.at(id);
				algn.filterOutAux(RGtag);
				algn.putAuxString("RG",replaceby);
			}
		}

		assert ( Aout[rgindex] );
		Aout[rgindex]->writeAlignment(algn);

		lineid += 1;

		if ( lineid % (1024*1024) == 0 )
		{
			std::cerr << "[V] " << lineid << std::endl;
		}
	}

	return EXIT_SUCCESS;
}

int main(int argc, char * argv[])
{
	try
	{
		std::vector<libmaus2::util::ArgParser::ArgumentDefinition> const Vformatin = libmaus2::bambam::BamAlignmentDecoderInfo::getArgumentDefinitions();
		std::vector<libmaus2::util::ArgParser::ArgumentDefinition> const Vformatout = libmaus2::bambam::BamBlockWriterBaseFactory::getArgumentDefinitions();

		std::vector<libmaus2::util::ArgParser::ArgumentDefinition> Vformat = libmaus2::util::ArgParser::mergeFormat(Vformatin,Vformatout);
		Vformat.push_back(libmaus2::util::ArgParser::ArgumentDefinition("","rgreplacetable",true));

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

		if ( arg.size() < 1 )
		{
			std::cerr << "[E] usage: " << arg.progname << " <outprefix>" << std::endl;
			return EXIT_FAILURE;
		}

		int const r = bamexploderg(arg);

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