/****
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 <inttypes.h>
#include <algorithm>
#include <atomic>
#include <exception>
#include <limits>
#include <memory>
#include <unordered_map>
#include "multinode.h"
#include "util/data_structures/queue.h"
#include "util/parallel/simple_thread_pool.h"
#include "util/log_stream.h"

using std::ifstream;
using std::unique_ptr;
using std::runtime_error;
using std::ofstream;
using std::string;
using std::vector;
using std::endl;
using std::tie;
using std::atomic;

struct OutputPart {
	size_t file;
	string data;
	string accs;
};

struct OutputChunk {
	static constexpr size_t POISON = std::numeric_limits<size_t>::max();
	size_t order = POISON;
	vector<OutputPart> parts;

	bool operator==(const OutputChunk& rhs) const {
		return order == rhs.order;
	}
};

static void write_blocks(Job& job, VolumedFile& volumes, vector<unique_ptr<ofstream>>& out, vector<unique_ptr<ofstream>>& acc_out, const vector<int>& block_mapping) {
	job.log("Writing length sorted blocks");
	OId oid = 0;
	atomic<uint64_t> bytes_written{ 0 };
	const size_t writer_count = std::max<size_t>(1, std::min<size_t>(out.size(), std::max(1, config.threads_)));
	const int formatter_count = std::max(1, config.threads_);
	SimpleThreadPool pool;

	auto process_volume = [&](const Volume& volume, ptrdiff_t idx) {
		const SequenceFile::Flags flags = SequenceFile::Flags::ALL | SequenceFile::Flags::NEED_LETTER_COUNT;
		unique_ptr<SequenceFile> file;
		try {
			file.reset(SequenceFile::auto_create({ volume.path }, flags, amino_acid_traits));
		} catch(FormatDetectionError& e) {
			throw runtime_error("Error opening file " + volume.path + ": " + e.what());
		}
		const uint64_t limit = std::min<uint64_t>(Util::String::interpret_number(config.memory_limit.get(DEFAULT_MEMORY_LIMIT)), SequenceFile::DEFAULT_LOAD_SIZE);
		uint64_t bytes = 0, ms = 0;
		for (;;) {
			TaskTimer timer;
			Block* b = file->load_seqs(limit);
			*message_stream << std::setprecision(2) << std::fixed << file->rel_file_ptr() * 100 << "%" << endl;
			ms += timer.microseconds();
			bytes += b->raw_bytes();
			if (b->empty()) {
				delete b;
				break;
			}
			const OId oid_begin = oid;
			const size_t seq_count = b->seqs().size();			
			const size_t target_chunks = std::max<size_t>(1, (size_t)formatter_count * 16);
			const size_t chunk_size = std::max<size_t>(1, std::min<size_t>(seq_count / target_chunks + 1, (size_t)8192));
			const size_t chunk_count = (seq_count + chunk_size - 1) / chunk_size;
			
			vector<unique_ptr<Queue<OutputChunk>>> queues;
			queues.reserve(writer_count);
			for (size_t i = 0; i < writer_count; ++i)
				queues.emplace_back(new Queue<OutputChunk>(std::max<size_t>(16, formatter_count * 4), formatter_count, 1, OutputChunk()));

			auto writer = [&](const atomic<bool>& stop, size_t writer_id) {
				std::unordered_map<size_t, OutputChunk> pending;
				size_t next_expected = 0;
				OutputChunk chunk;
				while (!stop.load(std::memory_order_relaxed) && queues[writer_id]->wait_and_dequeue(chunk)) {
					pending.emplace(chunk.order, std::move(chunk));
					auto it = pending.find(next_expected);
					while (it != pending.end()) {
						for (OutputPart& p : it->second.parts) {
							out[p.file]->write(p.data.data(), p.data.size());
							acc_out[p.file]->write(p.accs.data(), p.accs.size());
							bytes_written.fetch_add(p.data.size() + p.accs.size(), std::memory_order_relaxed);
						}
						pending.erase(it);
						++next_expected;
						it = pending.find(next_expected);
					}
				}
			};

			atomic<size_t> next(0);
			auto worker = [&](const atomic<bool>& stop) {
				vector<TextBuffer> buffers(out.size());
				vector<TextBuffer> acc_buffers(out.size());
				size_t c;
				while (!stop.load(std::memory_order_relaxed) && (c = next.fetch_add(1, std::memory_order_relaxed), c < chunk_count)) {
					const size_t j_begin = c * chunk_size;
					const size_t j_end = std::min(j_begin + chunk_size, seq_count);
					for (size_t j = j_begin; j < j_end; ++j) {
						const OId seq_oid = oid_begin + (OId)j;
						const size_t block = (size_t)block_mapping[seq_oid];
						const string oid_str = std::to_string(seq_oid);
						Util::Seq::format(b->seqs()[j], oid_str.c_str(), nullptr, buffers[block], "fasta", amino_acid_traits);
						acc_buffers[block] << oid_str << '\t';
						acc_buffers[block] << Util::Seq::seqid(b->ids()[j]) << '\n';
					}					
					vector<OutputChunk> bundles(writer_count);
					for (size_t w = 0; w < writer_count; ++w)
						bundles[w].order = c;
					for (size_t block = 0; block < buffers.size(); ++block) {
						if (buffers[block].size() == 0)
							continue;
						OutputPart part;
						part.file = block;
						part.data.assign(buffers[block].data(), buffers[block].size());
						part.accs.assign(acc_buffers[block].data(), acc_buffers[block].size());
						bundles[block % writer_count].parts.push_back(std::move(part));
						buffers[block].clear();
						acc_buffers[block].clear();
					}
					for (size_t w = 0; w < writer_count; ++w)
						queues[w]->enqueue(std::move(bundles[w]));
				}
				};

			vector<std::thread::id> writer_threads;
			for (size_t i = 0; i < writer_count; ++i)
				writer_threads.push_back(pool.spawn(writer, i));
			vector<std::thread::id> formatter_threads;
			for (int j = 0; j < formatter_count; ++j)
				formatter_threads.push_back(pool.spawn(worker));
			try {
				pool.join(formatter_threads.begin(), formatter_threads.end());
			}
			catch (...) {
				for (int i = 0; i < formatter_count; ++i)
					for (auto& q : queues)
						q->close();
				pool.join(writer_threads.begin(), writer_threads.end());
				delete b;
				throw;
			}			
			for (int i = 0; i < formatter_count; ++i)
				for (auto& q : queues)
					q->close();
			pool.join(writer_threads.begin(), writer_threads.end());
			oid += (OId)seq_count;
			delete b;
		}
		};

	TaskTimer timer;
	for (VolumedFile::const_iterator i = volumes.begin(); i != volumes.end(); ++i) {
		process_volume(*i, i - volumes.begin());
	}
	const int64_t t = timer.microseconds();
	const uint64_t total_bytes = bytes_written.load();
	job.log("Wrote %zu bytes to disk at %.2f MB/s", total_bytes, (double)total_bytes / MEGABYTES / (t / 1e6));
	for (const auto& file : out)
		if (!*file)
			throw runtime_error("Error writing length sorted block");
	for (const auto& file : acc_out)
		if (!*file)
			throw runtime_error("Error writing length sorted block");
}

string len_sort(Job& job, VolumedFile& volumes) {
	Atomic lock(job.root_dir() + "lensort_lock", job), done(job.root_dir() + "lensort_done", job);
	const bool first_round_linear = job.is_linear_round();
	const string input_parts = job.root_dir() + "input.tsv", input_letters = job.root_dir() + "input_letters.txt";
	if (lock.fetch_add() == 0) {
		ofstream idx(input_parts);
		vector<unique_ptr<ofstream>> out;
		vector<unique_ptr<ofstream>> acc_out;
		vector<int> block_mapping = make_blocks(job, volumes, out, acc_out);
		write_blocks(job, volumes, out, acc_out, block_mapping);
		done.fetch_add();
		job.finish_step();
	}
	else
		done.await(1);
	uint64_t letters, seq_count;
	ifstream input_letters_file(input_letters);
	input_letters_file >> letters >> seq_count;
	if(!input_letters_file)
		throw runtime_error("Error opening file " + input_letters);
	volumes.set_letter_count(letters);
	job.register_sync_file(input_letters);
	if (first_round_linear) {
		double block_gb;
		int index_chunks;
		tie(block_gb, index_chunks) = ::block_size(job.mem_limit, letters, Sensitivity::FAST, true, config.threads_); // TODO take cluster steps into account here
		config.lowmem_ = index_chunks;
	}
	job.set_max_oid(seq_count - 1);
	return input_parts;
}
