/****
DIAMOND protein sequence aligner
Copyright (C) 2012-2026 Benjamin J. Buchfink

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/>.
****/
// SPDX-License-Identifier: GPL-3.0-or-later

#include "basic/config.h"
#include "util/log_stream.h"
#include "masking/masking.h"
#include "data/taxonomy.h"
#include "data/taxon_list.h"
#include "data/taxonomy_nodes.h"
#include "murmurhash/MurmurHash3.h"
#include "legacy/dmnd/record_reader.h"
#include "util/parallel/multiprocessing.h"
#include "dmnd.h"
#include "util/system/system.h"
#include "util/algo/external_sort.h"
#include "util/util.h"
#include "data/fasta/fasta_file.h"
#include "util/sequence/sequence.h"
#include "legacy/dmnd/io.h"

using std::tuple;
using std::string;
using std::list;
using std::endl;
using std::unique_ptr;
using std::pair;
using std::runtime_error;
using std::vector;

const char* DatabaseFile::FILE_EXTENSION = ".dmnd";
const uint32_t ReferenceHeader::current_db_version_prot = 3;
const uint32_t ReferenceHeader::current_db_version_nucl = 4;

void serialize(File &s, const ReferenceHeader2 &h)
{
	s.write(big_endian_byteswap(sizeof(ReferenceHeader2)));
	s.write(h.hash, sizeof(h.hash));
	s.write(big_endian_byteswap(h.taxon_array_offset));
	s.write(big_endian_byteswap(h.taxon_array_size));
	s.write(big_endian_byteswap(h.taxon_nodes_offset));
	s.write(big_endian_byteswap(h.taxon_names_offset));
}

void deserialize(File &d, ReferenceHeader2 &h)
{
	DynamicRecordReader(d).read(h.hash, sizeof(h.hash))
		>> h.taxon_array_offset
		>> h.taxon_array_size
		>> h.taxon_nodes_offset
		>> h.taxon_names_offset
		>> Finish();
}

void deserialize(File& file, ReferenceHeader& h)
{
	file.read(h.magic_number);
	h.magic_number = big_endian_byteswap(h.magic_number);
	file.read(h.build);
	h.build = big_endian_byteswap(h.build);
	file.read(h.db_version);
	h.db_version = big_endian_byteswap(h.db_version);
	file.read(h.sequences);
	h.sequences = big_endian_byteswap(h.sequences);
	file.read(h.letters);
	h.letters = big_endian_byteswap(h.letters);
	file.read(h.pos_array_offset);
	h.pos_array_offset = big_endian_byteswap(h.pos_array_offset);
}

void serialize(File& file, const ReferenceHeader& h)
{
	file.write(big_endian_byteswap(h.magic_number));
	file.write(big_endian_byteswap(h.build));
	file.write(big_endian_byteswap(h.db_version));
	file.write(big_endian_byteswap(h.sequences));
	file.write(big_endian_byteswap(h.letters));
	file.write(big_endian_byteswap(h.pos_array_offset));
}

File& operator>>(File& file, SequenceFile::SeqInfo& r) {
	uint32_t p;
	file.read(r.pos);
	r.pos = big_endian_byteswap(r.pos);
	file.read(r.seq_len);
	r.seq_len = big_endian_byteswap(r.seq_len);
	file.read(p);
	return file;
}

void serialize(File& file, const SequenceFile::SeqInfo& r) {
	file.write(big_endian_byteswap(r.pos));
	file.write(big_endian_byteswap(r.seq_len));
	file.write((uint32_t)0);
}

SequenceFile::SeqInfo DatabaseFile::read_seqinfo() {
	SeqInfo r;
	file_.read(r);
	pos_array_offset += SeqInfo::SIZE;
	return r;
}

void DatabaseFile::putback_seqinfo() {
	pos_array_offset -= SeqInfo::SIZE;
}

void DatabaseFile::init(Flags flags)
{
	read_header(file_, ref_header);
	if (flag_any(flags, Flags::NO_COMPATIBILITY_CHECK))
		return;
	if (ref_header.build < min_build_required || ref_header.db_version < MIN_DB_VERSION)
		throw runtime_error("Database was built with an older version of Diamond and is incompatible.");
	if (ref_header.db_version > std::max(ReferenceHeader::current_db_version_prot , ReferenceHeader::current_db_version_nucl))
		throw runtime_error("Database was built with a newer version of Diamond and is incompatible.");
	if (ref_header.sequences == 0)
		throw runtime_error("Incomplete database file. Database building did not complete successfully.");
	deserialize(file_, header2);
	pos_array_offset = ref_header.pos_array_offset;
}

DatabaseFile::DatabaseFile(const string& input_file, Flags flags, const ValueTraits& value_traits) :
	SequenceFile(SequenceFile::Type::DMND, flags, FormatFlags::DICT_LENGTHS | FormatFlags::DICT_SEQIDS | FormatFlags::SEEKABLE | FormatFlags::LENGTH_LOOKUP, value_traits),
	file_(auto_append_extension_if_exists(input_file, FILE_EXTENSION), "rb")
{
	init(flags);

	vector<string> e;
	if (flag_any(flags, Flags::TAXON_MAPPING) && !has_taxon_id_lists())
		e.push_back("taxonomy mapping information (--taxonmap option)");
	if (flag_any(flags, Flags::TAXON_NODES) && !has_taxon_nodes())
		e.push_back("taxonomy nodes information (--taxonnodes option)");
	if (flag_any(flags, Flags::TAXON_SCIENTIFIC_NAMES) && !has_taxon_scientific_names())
		e.push_back("taxonomy names information (--taxonnames option)");
	if (flag_any(flags, Flags::TAXON_RANKS) && build_version() < 131)
		e.push_back("taxonomy ranks information (database needs to be built with diamond version >= 0.9.30");

	if (!e.empty())
		throw std::runtime_error("Options require taxonomy information included in the database. Please use the respective options to build this information into the database when running diamond makedb: " + join(", ", e.begin(), e.end()));

	if (flag_any(flags, Flags::TAXON_MAPPING)) {
		file_.seek(header2.taxon_array_offset);
		taxon_list_.reset(new TaxonList(file_, ref_header.sequences, header2.taxon_array_size));
	}

	if (flag_any(flags, Flags::TAXON_SCIENTIFIC_NAMES)) {
		file_.seek(header2.taxon_names_offset);
		deserialize(file_, taxon_scientific_names_);
	}

	if (flag_any(flags, Flags::TAXON_NODES)) {
		file_.seek(header2.taxon_nodes_offset);
		taxon_nodes_.reset(new TaxonomyNodes(file_, ref_header.build));
		init_cache();
	}

	if (flag_any(flags, Flags::ACC_TO_OID_MAPPING | Flags::OID_TO_ACC_MAPPING | Flags::NEED_LENGTH_LOOKUP))
		read_seqid_list();
}

/*DatabaseFile::DatabaseFile(TempFile& tmp_file, const ValueTraits& value_traits) :
	SequenceFile(SequenceFile::Type::DMND, Flags::NONE, FormatFlags::DICT_LENGTHS | FormatFlags::DICT_SEQIDS | FormatFlags::SEEKABLE | FormatFlags::LENGTH_LOOKUP, value_traits),
	InputFile(tmp_file, 0),
	temporary(true)
{
	init();
}*/

int64_t DatabaseFile::file_count() const {
	return 1;
}

void DatabaseFile::close() {
	/*if (temporary)
		InputFile::close_and_delete();
	else*/
	file_.close();
}

void DatabaseFile::read_header(File &stream, ReferenceHeader &header)
{
	deserialize(stream, header);
	if (header.magic_number != ReferenceHeader().magic_number)
		throw DatabaseFormatException();
}

bool DatabaseFile::has_taxon_id_lists() const
{
	return header2.taxon_array_offset != 0;
}

bool DatabaseFile::has_taxon_nodes() const
{
	return header2.taxon_nodes_offset != 0;
}

bool DatabaseFile::has_taxon_scientific_names() const {
	return header2.taxon_names_offset != 0;
}

static void push_seq(const Sequence &seq, const char *id, size_t id_len, uint64_t &offset, vector<SequenceFile::SeqInfo> &pos_array, File &out, size_t &letters, size_t &n_seqs)
{
	pos_array.emplace_back(offset, seq.length());
	out.write("\xff", 1);
	out.write(seq.data(), seq.length());
	out.write("\xff", 1);
	out.write(id, id_len + 1);
	letters += seq.length();
	++n_seqs;
	offset += seq.length() + id_len + 3;
}

void DatabaseFile::make_db()
{
	config.file_buffer_size = 4 * MEGABYTES;
	if (config.input_ref_file.size() > 1)
		throw runtime_error("Too many arguments provided for option --in.");
	const string input_file_name = config.input_ref_file.empty() ? string() : config.input_ref_file.front();
	if (input_file_name.empty())
		*message_stream << "Input file parameter (--in) is missing. Input will be read from stdin." << endl;
	else
		*message_stream << "Database input file: " << input_file_name << endl;

	TaskTimer total;
	TaskTimer timer("Opening the database file", true);

	value_traits = (config.dbtype == SequenceType::amino_acid) ? amino_acid_traits : nucleotide_traits;
    FastaFile db_file({ input_file_name }, Flags::NONE, value_traits);

    unique_ptr<File> out(new File(config.database, "wb"));
	ReferenceHeader header;
    ReferenceHeader2 header2;

	serialize(*out, header);
	serialize(*out, header2);

	size_t letters = 0, n = 0, n_seqs = 0, total_seqs = 0;
	uint64_t offset = out->tell();

	db_file.flags() |= SequenceFile::Flags::ALL;
    if (config.dbtype == SequenceType::nucleotide){
        header.db_version = ReferenceHeader::current_db_version_nucl;
        db_file.flags() |= SequenceFile::Flags::DNA_PRESERVATION;
    }

    Block* block;
	vector<SeqInfo> pos_array;
	ExternalSorter<pair<string, OId>> accessions;
	Util::Seq::AccessionParsing acc_stats;
	try {
		while (true) {
			timer.go("Loading sequences");
			block = db_file.load_seqs((int64_t)1e9, 0, nullptr);
			if (block->empty()) {
				delete block;
				break;
			}
			n = block->seqs().size();

			if (config.dbtype == SequenceType::amino_acid && config.masking_ != "0") {
				timer.go("Masking sequences");
				mask_seqs(block->seqs(), Masking::get(), false, MaskingAlgo::SEG);
			}

            timer.go("Writing sequences");
			for (size_t i = 0; i < n; ++i) {
				Sequence seq = block->seqs()[i];
				if (seq.length() == 0)
					throw std::runtime_error("File format error: sequence of length 0");
				push_seq(seq, block->ids()[i], block->ids().length(i), offset, pos_array, *out, letters, n_seqs);
			}
			if (!config.prot_accession2taxid.empty()) {
				timer.go("Writing accessions");
				for (size_t i = 0; i < n; ++i) {
					vector<string> acc = Util::Seq::accession_from_title(block->ids()[i], !config.no_parse_seqids, acc_stats);
					for (const string& s : acc)
						accessions.push(std::make_pair(s, total_seqs + i));
				}
			}
			timer.go("Hashing sequences");
			for (size_t i = 0; i < n; ++i) {
				Sequence seq = block->seqs()[i];
				MurmurHash3_x64_128(seq.data(), (int)seq.length(), header2.hash, header2.hash);
				MurmurHash3_x64_128(block->ids()[i], block->ids().length(i), header2.hash, header2.hash);
			}
			delete block;
			total_seqs += n;
		}
	}
	catch (std::exception&) {
		out->close();
		remove(out->name().c_str());
		throw;
	}

	timer.finish();

	timer.go("Writing trailer");
	header.pos_array_offset = offset;
	pos_array.emplace_back(offset, 0);
	for (const SeqInfo& r : pos_array)
		serialize(*out, r);
	pos_array.clear();
	pos_array.shrink_to_fit();
	timer.finish();

	if (!config.prot_accession2taxid.empty() && !config.no_parse_seqids)
		*message_stream << endl << "Accession parsing rules triggered for database seqids (use --no-parse-seqids to disable):" << endl << acc_stats << endl;

	Util::Table stats;
	stats("Database sequences", n_seqs);
	stats("Database letters", letters);

	taxonomy.init();
	if (!config.prot_accession2taxid.empty()) {
		header2.taxon_array_offset = out->tell();
		TaxonList::build(*out, accessions, n_seqs, stats);
		header2.taxon_array_size = out->tell() - header2.taxon_array_offset;
	}
	if (!config.nodesdmp.empty()) {
		TaxonomyNodes nodes(config.nodesdmp);
		header2.taxon_nodes_offset = out->tell();
		nodes.save(*out);
	}
	if (!config.namesdmp.empty()) {
		header2.taxon_names_offset = out->tell();
		serialize(*out, taxonomy.name_);
	}

#ifdef EXTRA
    header2.db_type = config.dbtype;
#endif

    timer.go("Closing the input file");
	db_file.close();

	timer.go("Closing the database file");
	header.letters = letters;
	header.sequences = n_seqs;
	out->seek(0);
	serialize(*out, header);
	serialize(*out, header2);
	out->close();

	timer.finish();
	stats("Database hash", hex_print(header2.hash, 16));
	stats("Total time", total.get(), "s");

	*message_stream << endl << stats;
}

void DatabaseFile::set_seqinfo_ptr(OId i) {
	pos_array_offset = ref_header.pos_array_offset + SeqInfo::SIZE * i;
}

OId DatabaseFile::tell_seq() const {
	return (pos_array_offset - ref_header.pos_array_offset) / SeqInfo::SIZE;
}

bool DatabaseFile::eof() const {
	return tell_seq() == sequence_count().value();
}

void DatabaseFile::init_seq_access() {
	file_.seek(sizeof(ReferenceHeader) + sizeof(ReferenceHeader2) + 8);
}

bool DatabaseFile::read_seq(vector<Letter>& seq, string &id, std::vector<char>* quals)
{
	char c;
	file_.read(c);
	seq.clear();
	id.clear();
	file_.read_to(std::back_inserter(seq), '\xff');
	file_.read_to(std::back_inserter(id), '\0');
	return false;
}

void DatabaseFile::skip_seq()
{
	char c;
	if(file_.read_max(&c, 1) != 1)
		throw std::runtime_error("Unexpected end of file.");
	file_.getdelim('\xff');
	file_.getdelim('\0');
}

bool DatabaseFile::is_diamond_db(const string &file_name) {
	if (file_name == "-" || file_name.empty())
		return false;
	File db_file(file_name, "rb");
	uint64_t magic_number = 0;
	db_file.read_max(&magic_number, sizeof(uint64_t));
	bool r = (magic_number == ReferenceHeader::MAGIC_NUMBER);
	db_file.close();
	return r;
}

void DatabaseFile::create_partition_fixednumber(size_t n) {
	size_t max_letters_balanced = static_cast<size_t>(std::ceil(static_cast<double>(ref_header.letters)/static_cast<double>(n)));
	*message_stream << "Fixed number partitioning using " << max_letters_balanced << " (" << n << ")" << endl;
	this->create_partition(max_letters_balanced);
}

void DatabaseFile::create_partition_balanced(int64_t max_letters) {
	//double n = std::ceil(static_cast<double>(ref_header.letters) / static_cast<double>(max_letters));
	//size_t max_letters_balanced = static_cast<size_t>(std::ceil(static_cast<double>(ref_header.letters)/n));
	//cout << "Balanced partitioning using " << max_letters_balanced << " (" << max_letters << ")" << endl;
	this->create_partition(max_letters);
}

void DatabaseFile::create_partition(size_t max_letters) {
	TaskTimer timer("Create partition of DatabaseFile");
	size_t letters = 0, seqs = 0, total_seqs = 0;
	int i_chunk = 0;

	size_t oid = 0, oid_begin;
	set_seqinfo_ptr(oid);
	init_seqinfo_access();

	SeqInfo r = read_seqinfo();
	bool first = true;

	while (r.seq_len) {
		SeqInfo r_next = read_seqinfo();
		if (first) {
			oid_begin = oid;
			first = false;
		}
		letters += r.seq_len;
		++seqs;
		++total_seqs;
		if ((letters > max_letters) || (r_next.seq_len == 0)) {
			partition.chunks.push_back(Chunk(i_chunk, oid_begin, seqs));
			first = true;
			seqs = 0;
			letters = 0;
			++i_chunk;
		}
		r = r_next;
		++oid;
	}

	reverse(partition.chunks.begin(), partition.chunks.end());

	partition.max_letters = max_letters;
	partition.n_seqs_total = total_seqs;
}

int DatabaseFile::get_n_partition_chunks() {
	return (int)partition.chunks.size();
}

void DatabaseFile::save_partition(const string & partition_file_name, const string & annotation) {
	std::ofstream out(partition_file_name);
	// cout << "WRITING " << partition_file_name << endl;
	for (const auto& i : partition.chunks) {
		out << to_string(i);
		if (annotation.size() > 0) {
			out << " " << annotation;
		}
		out << endl;
	}
}

Chunk to_chunk(const string & line) {
	vector<string> tokens = split(line, ' ');
	return Chunk((int)stoull(tokens[0]), stoull(tokens[1]), stoull(tokens[2]));
}

string to_string(const Chunk & c) {
	const string buf = std::to_string(c.i) + " " + std::to_string(c.offset) + " " + std::to_string(c.n_seqs);
	return buf;
}

void DatabaseFile::load_partition(const string & partition_file_name) {
	string line;
	std::ifstream in(partition_file_name);
	clear_partition();
	while (getline(in, line)) {
		auto chunk = to_chunk(line);
		partition.chunks.push_back(chunk);
	}
}

void DatabaseFile::clear_partition() {
	partition.max_letters = 0;
	partition.n_seqs_total = 0;
	partition.chunks.clear();
}

void DatabaseFile::init_seqinfo_access() {
	file_.seek(pos_array_offset);
}

void DatabaseFile::seek_chunk(const Chunk& chunk) {
	set_seqinfo_ptr(chunk.offset);
}

size_t DatabaseFile::id_len(const SeqInfo& seq_info, const SeqInfo& seq_info_next) {
	return seq_info_next.pos - seq_info.pos - seq_info.seq_len - 3;
}

void DatabaseFile::seek_offset(size_t p) {
	file_.seek(p);
}

void DatabaseFile::read_seq_data(Letter* dst, size_t len, size_t& pos, bool seek) {
	if (seek)
		file_.seek(pos);
	file_.read(dst - 1, len + 2);
	*(dst - 1) = Sequence::DELIMITER;
	*(dst + len) = Sequence::DELIMITER;
}

void DatabaseFile::read_id_data(const int64_t oid, char* dst, size_t len, bool all, bool full_titles) {
	file_.read(dst, len + 1);
}

void DatabaseFile::skip_id_data() {
	file_.getdelim('\0');
}

optional<uint64_t> DatabaseFile::sequence_count() const {
	return ref_header.sequences;
}

optional<uint64_t> DatabaseFile::letters() const {
	return ref_header.letters;
}

int DatabaseFile::db_version() const {
	return ref_header.db_version;
}

int DatabaseFile::program_build_version() const {
	return ref_header.build;
}

int DatabaseFile::build_version() {
	return ref_header.build;
}

int DatabaseFile::rank(TaxId taxid) const
{
	const int r = taxon_nodes_->rank(taxid);
	return ref_header.build <= 165 && r >= 0 ? Rank::legacy_rank_codes.at(r) : r;
}

DatabaseFile::~DatabaseFile()
{
	close();
}

DbFilter* DatabaseFile::filter_by_accession(const std::string& file_name)
{
	throw std::runtime_error("The .dmnd database format does not support filtering by accession.");
	return nullptr;
}

std::string DatabaseFile::file_name()
{
	return file_.name();
}

std::vector<TaxId> DatabaseFile::taxids(size_t oid) const
{
	return (*taxon_list_)[oid];
}

void DatabaseFile::seq_data(size_t oid, std::vector<Letter>& dst)
{
	throw OperationNotSupported();
}

Loc DatabaseFile::seq_length(size_t oid)
{
	if (oid < seq_length_.size())
		return seq_length_[oid];
	throw std::out_of_range("DatabaseFile::seq_length");
}

void DatabaseFile::end_random_access(bool dictionary)
{
	if (!dictionary)
		return;
	free_dictionary();
}

string DatabaseFile::taxon_scientific_name(TaxId taxid) const {
	return taxid < (TaxId)taxon_scientific_names_.size() && !taxon_scientific_names_[taxid].empty() ? taxon_scientific_names_[taxid] : std::to_string(taxid);
}

void DatabaseFile::read_seqid_list() {
	if (flag_any(flags_, Flags::ACC_TO_OID_MAPPING))
		acc2oid_.reserve(sequence_count().value());
	if (flag_any(flags_, Flags::NEED_LENGTH_LOOKUP))
		seq_length_.reserve(sequence_count().value());
	OId oid = 0;
	vector<Letter> seq;
	string id;
	init_seq_access();
	for (uint64_t n = 0; n < sequence_count().value(); ++n) {
		read_seq(seq, id);
		const char* msg = Util::Seq::fix_title(id);
		if (msg)
			*message_stream << "Warning: " << msg << std::endl;
		add_seqid_mapping(id, oid++);
		if (flag_any(flags_, Flags::NEED_LENGTH_LOOKUP))
			seq_length_.push_back((Loc)seq.size());
	}
}

TaxId DatabaseFile::max_taxid() const {
	return taxon_nodes_->max();
}

TaxId DatabaseFile::get_parent(TaxId taxid) {
	return taxon_nodes_->get_parent(taxid);
}

void DatabaseFile::print_info() const {
	*message_stream << "Database: " << config.database << ' ';
	*message_stream << "(type: " << to_string(type()) << ", ";
	*message_stream << "sequences: " << sequence_count().value() << ", ";
	*message_stream << "letters: " << letters().value() << ", build version: " << ref_header.build << ')' << endl;
}