/****
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

#ifndef WIN32
#include <sys/mman.h>
#include <fcntl.h>
#endif
#include "hit_buffer.h"
#include "basic/config.h"
#include "util/log_stream.h"
#include "util/parallel/simple_thread_pool.h"
#include "data/block/block.h"
#include "util/memory/mem_profile.h"

using std::vector;
using std::string;
using std::thread;
using std::copy;
using std::runtime_error;
using std::mutex;
using std::endl;
using std::atomic_size_t;
using std::pair;
using std::atomic;
using std::tuple;

namespace Search {

HitBuffer::HitBuffer(const vector<Key>& key_partition, const string& tmpdir, bool long_subject_offsets, int query_contexts, int thread_count, Config& cfg) :
	HitBuffer(key_partition, tmpdir, long_subject_offsets, query_contexts, thread_count,
		(uint32_t)cfg.query->seqs().size(), (uint64_t)cfg.target->seqs().raw_len(),
		cfg.search_pool, cfg.lin_index.get() != nullptr)
{}

HitBuffer::HitBuffer(const vector<Key>& key_partition, const string& tmpdir, bool long_subject_offsets, int query_contexts, int thread_count, uint32_t max_query, uint64_t max_target, SimpleThreadPool& search_pool, bool hash_key) :
	key_partition_(key_partition),
	hash_key_(hash_key),
	long_subject_offsets_(long_subject_offsets),
	query_contexts_(query_contexts),
	max_query_(max_query),
	max_target_(max_target),
	search_pool_(search_pool),
	bins_processed_(0),
	total_disk_size_(0),
	load_worker_(nullptr),
	load_exception_(nullptr),
	aborted_(false)
{
	*log_stream << "Async_buffer() " << key_partition.back() << std::endl;
	count_ = new atomic_size_t[key_partition.size()];

	for (size_t i = 0; i < key_partition.size(); ++i) {
		count_[i].store(0, std::memory_order_relaxed);
		if (config.swipe_all)
			continue;
		if (config.trace_pt_membuf) {
			membuf_out_queue_.push_back(new Queue<std::pair<int, std::vector<Hit>*>>(thread_count * 4, thread_count, 1, pair<int, vector<Hit>*>(0, nullptr)));
			hit_buf_.emplace_back();
			hit_buf_.back().reserve(1 * GIGABYTES / sizeof(Hit));
		}
		else {
			out_queue_.push_back(new Queue<tuple<int, TextBuffer*, uint32_t>>(thread_count * 4, thread_count, 1, tuple<int, TextBuffer*, uint32_t>(0, nullptr, 0)));
			tmp_file_.push_back(File(Temporary()));
		}
	}

	if (!config.swipe_all) {
		for (size_t i = 0; i < key_partition.size(); ++i) {
			writer_.push_back(search_pool_.spawn_method(this, &HitBuffer::write_worker, (int)i));
		}
	}
}

void HitBuffer::write_worker(const std::atomic<bool>& stop, int bin) {
	try {
		if (config.trace_pt_membuf) {
			while(!stop) {
				pair<int, vector<Hit>*> buf;
				if (!membuf_out_queue_[bin]->wait_and_dequeue(buf))
					break;
				vector<Hit>& hits = hit_buf_[buf.first];
				hits.insert(hits.end(), buf.second->begin(), buf.second->end());
				delete buf.second;
			}
		}
		else {
			while (!stop) {
				tuple<int, TextBuffer*, uint32_t> buf;
				if (!out_queue_[bin]->wait_and_dequeue(buf))
					break;
				File& tmp_file = tmp_file_[std::get<0>(buf)];
				tmp_file.write(std::get<1>(buf)->size());
				tmp_file.write(std::get<2>(buf));
				tmp_file.write(std::get<1>(buf)->data(), std::get<1>(buf)->size());
				delete std::get<1>(buf);
			}
		}
	}
	catch (...) {
		abort();
		throw;
	}
}

void HitBuffer::abort() noexcept {
	if (aborted_.exchange(true, std::memory_order_acq_rel))
		return;
	for (auto q : membuf_out_queue_)
		if (q)
			q->abort();
	for (auto q : out_queue_)
		if (q)
			q->abort();
}

void HitBuffer::finish_writing() {
	auto cleanup = [&]() {
		for (auto& q : membuf_out_queue_)
			delete q;
		for (auto& q : out_queue_)
			delete q;
		writer_.clear();
		membuf_out_queue_.clear();
		out_queue_.clear();
		};
	if (config.trace_pt_membuf) {
		for (int i = 0; i < membuf_out_queue_[0]->producer_count(); ++i)
			for (int j = 0; j < bins(); ++j)
				membuf_out_queue_[j]->close();
	}
	else
		for (int i = 0; i < out_queue_[0]->producer_count(); ++i)
			for (int j = 0; j < bins(); ++j) {
				out_queue_[j]->close();
			}
	try {
		search_pool_.join(writer_.begin(), writer_.end());
	}
	catch (...) {
		cleanup();
		throw;
	}
	cleanup();
}

HitBuffer::~HitBuffer() noexcept(false) {
	if (!writer_.empty())
		throw runtime_error("HitBuffer::~HitBuffer(): writer thread still active");
	delete[] count_;
}

void HitBuffer::plan_bin_groups(uint64_t max_bytes) {
	const int n = bins();
	bin_group_end_.assign(n, 0);
	const bool combine = !config.trace_pt_membuf && !config.swipe_all;
	int b = 0;
	while (b < n) {
		int e = b + 1;
		uint64_t size = (uint64_t)count_[b] * sizeof(Hit);
		if (combine)
			while (e < n && size + (uint64_t)count_[e] * sizeof(Hit) <= max_bytes) {
				size += (uint64_t)count_[e] * sizeof(Hit);
				++e;
			}
		for (int i = b; i < e; ++i)
			bin_group_end_[i] = e;
		b = e;
	}
}

bool HitBuffer::load() {
	if (load_worker_)
		throw runtime_error("HitBuffer::load(): previous load still in progress");
	data_size_next_ = 0;
	auto worker = [&](int end) {
		try {
			Hit* out = data_loading_;
			for (; bins_processed_ < end; ++bins_processed_) {
				load_bin(out, bins_processed_);
				out += count_[bins_processed_];
			}
		}
		catch (...) {
			this->load_exception_ = std::current_exception();
		}
		};
	if (bins_processed_ == bins()) {
		return false;
	}
	const int begin = bins_processed_, end = group_end(begin);
	if (!config.trace_pt_membuf && !config.swipe_all) {
		size_t size = 0, disk_size = 0;
		for (int i = begin; i < end; ++i) {
			size += count_[i];
			disk_size += tmp_file_[i].size();
		}
		*log_stream << "Async_buffer.load() bins=[" << begin << ',' << end << ") " << size << " (" << (double)size * sizeof(Hit) / (1 << 30) << " GB, " << (double)disk_size / (1 << 30) << " GB on disk)" << endl;
		total_disk_size_ += disk_size;
		data_size_next_ = size;
		load_worker_ = new thread(worker, end);
	}
	input_range_next_.first = this->begin(begin);
	input_range_next_.second = this->end(end - 1);	
	return true;
}

void HitBuffer::load_bin(Hit* out, int bin)
{
	using namespace std::placeholders;
	if (config.trace_pt_membuf)
		return;
#if !_MSC_VER && !__APPLE__
	if (bin < bins() - 1) {
		posix_fadvise(fileno(tmp_file_[bin + 1].file()), 0, 0, POSIX_FADV_SEQUENTIAL);
		posix_fadvise(fileno(tmp_file_[bin + 1].file()), 0, 0, POSIX_FADV_WILLNEED);
	}
#endif
	File& f = tmp_file_[bin];
	if (count_[bin] == 0) {
		f.close();
		return;
	}
	const int p = config.threads_ > 1 ? std::min(config.threads_ - 1, config.load_threads) : 1;
	Queue<pair<vector<char>*, uint32_t>> queue(p * 4, 1, p, pair<vector<char>*, uint32_t>(nullptr, 0));
	atomic<Hit*> out_ptr(out);
	atomic<uint64_t> count(0);
	f.seek(0, SEEK_SET);
	SimpleThreadPool pool;
	auto worker = [&](const atomic<bool>& stop, int worker_id) {
		uint64_t my_count = 0;
		while (true) {
			pair<vector<char>*, uint32_t> v;
			if (!queue.wait_and_dequeue(v)) {
				break;
			}
			Hit* dst = out_ptr.fetch_add(v.second, std::memory_order_relaxed);
			uint32_t package_count = 0;
			vector<char>::const_iterator ptr = v.first->begin(), end = v.first->end();
			uint16_t nullscore;
			memcpy(&nullscore, &*ptr, 2);
			ptr += 2;
			while(ptr < end) {
				uint32_t query_id, seed_offset;				
				memcpy(&query_id, &*ptr, 4);
				if (query_id >= max_query_)
					throw runtime_error("HitBuffer::load_bin(): invalid query id / possibly corrupted temporary file");
				ptr += 4;
				memcpy(&seed_offset, &*ptr, 4);
				ptr += 4;
				PackedLoc subject_loc;
				uint32_t x;
				while (ptr < end) {
					uint16_t score;
					memcpy(&score, &*ptr, 2);
					ptr += 2;
					if (score == 0)
						break;
					if (long_subject_offsets_) {
						memcpy(&subject_loc, &*ptr, sizeof(PackedLoc));
						ptr += sizeof(PackedLoc);
					}
					else {
						memcpy(&x, &*ptr, 4);
						ptr += 4;
						subject_loc = x;
					}
					if((uint64_t)subject_loc >= max_target_)
						throw runtime_error("HitBuffer::load_bin(): invalid subject location / possibly corrupted temporary file");
					if(package_count >= v.second)
						throw runtime_error("HitBuffer::load_bin(): buffer overflow / possibly corrupted temporary file");
#ifdef HIT_KEEP_TARGET_ID
					uint32_t target_block_id;
					memcpy(&target_block_id, &*ptr, 4);
					ptr += 4;
					dst->target_block_id = target_block_id;
#endif
					dst->query_ = query_id;
					dst->subject_ = subject_loc;
					dst->seed_offset_ = seed_offset;
					dst->score_ = score;
					++dst;
					++package_count;
				}
			}
			delete v.first;
			my_count += package_count;
		}
		count.fetch_add(my_count, std::memory_order_relaxed);
		};
	for (int i = 0; i < p; ++i)
		pool.spawn(worker, i);
	try {
		while (!pool.stop()) {
			uint32_t n;
			size_t l;
			if (f.read_max(&l, sizeof(size_t)) < sizeof(size_t))
				break;
			f.read(n);
			vector<char>* v = new vector<char>(l);
			try {
				f.read(v->data(), l);
			}
			catch (...) {
				delete v;
				throw;
			}
			queue.enqueue(pair<vector<char>*, uint32_t>(v, n));
		}
	}
	catch (...) {
		queue.close();
		try { pool.join_all(); } catch (...) {}
		throw;
	}
	queue.close();
	pool.join_all();

	if (count != count_[bin])
		throw runtime_error("Mismatching hit count / possibly corrupted temporary file");
	f.close();
}

void HitBuffer::alloc_buffer() {
	MEM_SCOPE("align/trace-point-buffer");
	mmap_finished_ = mmap_loading_ = false;
	if (config.trace_pt_membuf)
		return;
	int64_t max_size = 0;
	for (int b = 0; b < bins(); b = group_end(b))
		max_size = std::max(max_size, group_size(b));
	alloc_size_ = max_size;
	if (max_size == 0) {
		data_loading_ = data_finished_ = nullptr;
		return;
	}
#ifdef _MSC_VER
	data_loading_ = new Hit[max_size];
	data_finished_ = new Hit[max_size];	
#else
	int flags = MAP_PRIVATE | MAP_ANONYMOUS;
#ifdef MAP_HUGETLB
	flags |= MAP_HUGETLB;
#endif
	mmap_finished_ = mmap_loading_ = true;
	data_finished_ = (Hit*)mmap(nullptr, max_size * sizeof(Hit), PROT_READ | PROT_WRITE, flags, -1, 0);
	data_loading_ = (Hit*)mmap(nullptr, max_size * sizeof(Hit), PROT_READ | PROT_WRITE, flags, -1, 0);
	if (data_finished_ == MAP_FAILED) {
		data_finished_ = new Hit[max_size];
		mmap_finished_ = false;
	}
	if (data_loading_ == MAP_FAILED) {
		data_loading_ = new Hit[max_size];
		mmap_loading_ = false;
	}
	if (mmap_finished_)
		MEM_ADD("align/trace-point-buffer", max_size * (int64_t)sizeof(Hit));
	if (mmap_loading_)
		MEM_ADD("align/trace-point-buffer", max_size * (int64_t)sizeof(Hit));
#endif
}

void HitBuffer::free_buffer() {
	if (!config.trace_pt_membuf) {
#ifdef _MSC_VER
		delete[] data_loading_;
		delete[] data_finished_;
#else
		if (mmap_finished_) {
			MEM_SUB("align/trace-point-buffer", alloc_size_ * (int64_t)sizeof(Search::Hit));
			munmap(data_finished_, alloc_size_ * sizeof(Search::Hit));
		}
		else
			delete[] data_finished_;
		if (mmap_loading_) {
			MEM_SUB("align/trace-point-buffer", alloc_size_ * (int64_t)sizeof(Search::Hit));
			munmap(data_loading_, alloc_size_ * sizeof(Search::Hit));
		}
		else
			delete[] data_loading_;
#endif
	}

}

}
