#ifndef _CHECKPOINT_H_
#define _CHECKPOINT_H_

#include "BloomDBG/AssemblyCounters.h"
#include "BloomDBG/AssemblyParams.h"
#include "BloomDBG/AssemblyStreams.h"
#include "Common/IOUtil.h"
#include "DataLayer/FastaReader.h"
#include <cstdio>
#include <fstream>
#include <iostream>

namespace BloomDBG {
const std::string CHECKPOINT_FASTA_EXT = ".contigs.fa";
const std::string CHECKPOINT_COUNTERS_EXT = ".counters.tsv";
const std::string CHECKPOINT_BLOOM_DBG_EXT = ".dbg.bloom";
const std::string CHECKPOINT_BLOOM_VISITED_EXT = ".visited.bloom";
const std::string CHECKPOINT_TMP_EXT = ".tmp";

/**
 * Save the current assembly state to a set of checkpoint files,
 * so that the assembly can be resumed from that point.
 *
 * @param dbg Bloom filter de Bruijn graph
 * @param visitedKmerSet k-mers included in output contigs so far
 * @param counters counters capturing assembly state
 * (e.g. next input read index, next output contig index)
 * @param params various assembly parameters (corresponding to
 * command line opts)
 */
template<typename BloomDBGT, typename VisitedKmerSetT>
static inline void
createCheckpoint(
    const BloomDBGT& dbg,
    const VisitedKmerSetT& visitedKmerSet,
    const AssemblyCounters& counters,
    const AssemblyParams& params)
{
	assert(params.checkpointsEnabled());
	assert(!params.checkpointPathPrefix.empty());
	std::string prefix = params.checkpointPathPrefix;

	if (params.verbose)
		std::cerr << "Writing checkpoint data..." << std::endl;

	/* write out Bloom filter de Bruijn graph */

	std::string dbgPath = prefix + CHECKPOINT_BLOOM_DBG_EXT;
	std::string dbgPathTmp = dbgPath + CHECKPOINT_TMP_EXT;

	if (params.verbose)
		std::cerr << '\t' << "Writing Bloom filter de Bruijn graph to `" << dbgPathTmp << "'"
		          << std::endl;

	std::ofstream dbgOut;
	dbgOut.open(dbgPathTmp.c_str());
	assert_good(dbgOut, dbgPathTmp);
	dbgOut << dbg;
	assert_good(dbgOut, dbgPathTmp);

	/* write out visited k-mers Bloom filter */

	std::string visitedPath = prefix + CHECKPOINT_BLOOM_VISITED_EXT;
	std::string visitedPathTmp = visitedPath + CHECKPOINT_TMP_EXT;

	if (params.verbose)
		std::cerr << '\t' << "Writing visited k-mers Bloom to `" << visitedPathTmp << "'"
		          << std::endl;

	std::ofstream visitedOut;
	visitedOut.open(visitedPathTmp.c_str());
	assert_good(visitedOut, visitedPathTmp);
	visitedOut << visitedKmerSet;
	assert_good(visitedOut, visitedPathTmp);

	/* write out index of next input read */

	std::string countersPath = prefix + CHECKPOINT_COUNTERS_EXT;
	std::string countersPathTmp = countersPath + CHECKPOINT_TMP_EXT;

	if (params.verbose)
		std::cerr << '\t' << "Writing assembly counters to `" << countersPathTmp << "'"
		          << std::endl;

	std::ofstream countersOut;
	countersOut.open(countersPathTmp.c_str());
	assert_good(countersOut, countersPathTmp);
	countersOut << counters;
	assert_good(countersOut, countersPathTmp);

	/* copy/move new checkpoint files on top of old ones */

	std::string fastaPath = prefix + CHECKPOINT_FASTA_EXT;
	std::string fastaPathTmp = fastaPath + CHECKPOINT_TMP_EXT;

	if (params.verbose)
		std::cerr << '\t' << "Copying `" << fastaPathTmp << " to `" << fastaPath << "'"
		          << std::endl;

	copyFile(fastaPathTmp, fastaPath);

	if (params.verbose)
		std::cerr << '\t' << "Moving `" << dbgPathTmp << "' to `" << dbgPath << "'" << std::endl;

	if (rename(dbgPathTmp.c_str(), dbgPath.c_str()) != 0) {
		perror("Error renaming file");
		abort();
	}

	if (params.verbose)
		std::cerr << '\t' << "Moving `" << visitedPathTmp << "' to `" << visitedPath << "'"
		          << std::endl;

	if (rename(visitedPathTmp.c_str(), visitedPath.c_str()) != 0) {
		perror("Error renaming file");
		abort();
	}

	if (params.verbose)
		std::cerr << '\t' << "Moving `" << countersPathTmp << "' to `" << countersPath << "'"
		          << std::endl;

	if (rename(countersPathTmp.c_str(), countersPath.c_str()) != 0) {
		perror("Error renaming file");
		abort();
	}
}

/** Return true if checkpoint files exist and are readable */
static inline bool
checkpointExists(const AssemblyParams& params)
{
	assert(params.checkpointsEnabled());
	assert(!params.checkpointPathPrefix.empty());
	std::string prefix = params.checkpointPathPrefix;

	std::string fastaPath = prefix + CHECKPOINT_FASTA_EXT;
	std::string dbgPath = prefix + CHECKPOINT_BLOOM_DBG_EXT;
	std::string visitedPath = prefix + CHECKPOINT_BLOOM_VISITED_EXT;
	std::string countersPath = prefix + CHECKPOINT_COUNTERS_EXT;

	return std::ifstream(fastaPath.c_str()).good() && std::ifstream(dbgPath.c_str()).good() &&
	       std::ifstream(visitedPath.c_str()).good() && std::ifstream(countersPath.c_str()).good();
}

/**
 * Restore assembly state from checkpoint files.
 *
 * @param dbg Bloom filter de Bruijn graph
 * @param visitedKmerSet k-mers included in output contigs so far
 * @param counters counters capturing assembly state
 * (e.g. next input read index, next output contig index)
 * @param params various assembly parameters (corresponding to
 * @param out main output stream for assembled contigs
 * command line opts)
 */
template<typename BloomDBGT, typename VisitedKmerSetT, typename InputReadStreamT>
static inline void
resumeFromCheckpoint(
    BloomDBGT& dbg,
    VisitedKmerSetT& visitedKmerSet,
    AssemblyCounters& counters,
    const AssemblyParams& params,
    AssemblyStreams<InputReadStreamT>& streams)
{
	assert(checkpointExists(params));
	assert(params.checkpointsEnabled());
	assert(!params.checkpointPathPrefix.empty());
	std::string prefix = params.checkpointPathPrefix;

	if (params.verbose)
		std::cerr << "Resuming from last checkpoint..." << std::endl;

	/* load Bloom filter de Bruijn graph */

	std::string dbgPath = prefix + CHECKPOINT_BLOOM_DBG_EXT;
	if (params.verbose)
		std::cerr << '\t' << "Reading Bloom filter de Bruijn graph from `" << dbgPath << "'"
		          << std::endl;
	dbg.loadFilter(dbgPath);

	/* load visited k-mers Bloom filter */

	std::string visitedPath = prefix + CHECKPOINT_BLOOM_VISITED_EXT;
	if (params.verbose)
		std::cerr << '\t' << "Reading reading visited k-mers Bloom from `" << visitedPath << "'"
		          << std::endl;
	visitedKmerSet.loadFilter(visitedPath);

	/* load index for next input read */

	std::string countersPath = prefix + CHECKPOINT_COUNTERS_EXT;
	if (params.verbose)
		std::cerr << '\t' << "Reading index of next input read from `" << countersPath << "'"
		          << std::endl;
	std::ifstream countersIn(countersPath.c_str());
	assert_good(countersIn, countersPath);
	countersIn >> counters;
	assert_good(countersIn, countersPath);

	/* advance to previous position in input reads */

	if (params.verbose)
		std::cerr << '\t' << "Advancing to read index " << counters.readsProcessed
		          << " in input reads..." << std::endl;

	FastaRecord rec;
	for (size_t i = 0; i < counters.readsProcessed && streams.in; ++i)
		streams.in >> rec;
	assert(!streams.in.eof());

	/* restore previously assembled contigs */

	std::string fastaPath = prefix + CHECKPOINT_FASTA_EXT;
	std::string fastaPathTmp = fastaPath + CHECKPOINT_TMP_EXT;
	if (params.verbose)
		std::cerr << '\t' << "Copying `" << fastaPath << "' to `" << fastaPathTmp << "'"
		          << std::endl;
	copyFile(fastaPath, fastaPathTmp);

	if (params.verbose)
		std::cerr << '\t' << "Outputting previously assembled contigs "
		          << "from `" << fastaPath << "'" << std::endl;
	std::ifstream prevContigs(fastaPath.c_str());
	assert_good(prevContigs, fastaPath);
	streams.out << prevContigs.rdbuf();
	assert(streams.out);
}

/** Delete a file if it exists */
static inline void
removeFileIfExists(const std::string& path)
{
	if (std::ifstream(path.c_str()).good()) {
		if (remove(path.c_str()) != 0) {
			perror("Error removing file");
			abort();
		}
	}
}

/** Remove all checkpoint-related files */
static inline void
removeCheckpointData(const AssemblyParams& params)
{
	assert(params.checkpointsEnabled());
	assert(!params.checkpointPathPrefix.empty());
	std::string prefix = params.checkpointPathPrefix;

	if (params.verbose)
		std::cerr << "Removing checkpoint files..." << std::endl;

	/* remove Bloom filter de Bruijn graph file(s) */

	std::string dbgPath = prefix + CHECKPOINT_BLOOM_DBG_EXT;
	std::string dbgPathTmp = dbgPath + CHECKPOINT_TMP_EXT;
	removeFileIfExists(dbgPath);
	removeFileIfExists(dbgPathTmp);

	/* remove visited k-mers Bloom filter file(s) */

	std::string visitedPath = prefix + CHECKPOINT_BLOOM_VISITED_EXT;
	std::string visitedPathTmp = visitedPath + CHECKPOINT_TMP_EXT;
	removeFileIfExists(visitedPath);
	removeFileIfExists(visitedPathTmp);

	/* remove assembly counters file(s) */

	std::string countersPath = prefix + CHECKPOINT_COUNTERS_EXT;
	std::string countersPathTmp = countersPath + CHECKPOINT_TMP_EXT;
	removeFileIfExists(countersPath);
	removeFileIfExists(countersPathTmp);

	/* remove contigs FASTA file(s) */

	std::string fastaPath = prefix + CHECKPOINT_FASTA_EXT;
	std::string fastaPathTmp = fastaPath + CHECKPOINT_TMP_EXT;
	removeFileIfExists(fastaPath);
	removeFileIfExists(fastaPathTmp);
}
}

#endif
