#include <utility>

#include "duckdb/common/arrow/arrow_util.hpp"
#include "duckdb/common/arrow/arrow_appender.hpp"
#include "duckdb/common/types/data_chunk.hpp"
#include "duckdb/function/table/arrow/arrow_duck_schema.hpp"
namespace duckdb {

bool ArrowUtil::TryFetchChunk(ChunkScanState &scan_state, ClientProperties options, idx_t batch_size, ArrowArray *out,
                              idx_t &count, ErrorData &error,
                              unordered_map<idx_t, const shared_ptr<ArrowTypeExtensionData>> extension_type_cast) {
	count = 0;
	ArrowAppender appender(scan_state.Types(), batch_size, std::move(options), std::move(extension_type_cast));
	const auto remaining_tuples_in_chunk = scan_state.RemainingInChunk();
	if (remaining_tuples_in_chunk) {
		// We start by scanning the non-finished current chunk
		idx_t cur_consumption = MinValue(remaining_tuples_in_chunk, batch_size);
		count += cur_consumption;
		auto &current_chunk = scan_state.CurrentChunk();
		appender.Append(current_chunk, scan_state.CurrentOffset(), scan_state.CurrentOffset() + cur_consumption,
		                current_chunk.size());
		scan_state.IncreaseOffset(cur_consumption);
	}
	while (count < batch_size) {
		if (!scan_state.LoadNextChunk(error)) {
			if (scan_state.HasError()) {
				error = scan_state.GetError();
			}
			return false;
		}
		if (scan_state.ChunkIsEmpty()) {
			// The scan was successful, but an empty chunk was returned
			break;
		}
		auto &current_chunk = scan_state.CurrentChunk();
		if (scan_state.Finished() || current_chunk.size() == 0) {
			break;
		}
		// The amount we still need to append into this chunk
		auto remaining = batch_size - count;

		// The amount remaining, capped by the amount left in the current chunk
		auto to_append_to_batch = MinValue(remaining, scan_state.RemainingInChunk());
		appender.Append(current_chunk, 0, to_append_to_batch, current_chunk.size());
		count += to_append_to_batch;
		scan_state.IncreaseOffset(to_append_to_batch);
	}
	if (count > 0) {
		*out = appender.Finalize();
	}
	return true;
}

idx_t ArrowUtil::FetchChunk(ChunkScanState &scan_state, ClientProperties options, idx_t chunk_size, ArrowArray *out,
                            const unordered_map<idx_t, const shared_ptr<ArrowTypeExtensionData>> &extension_type_cast) {
	ErrorData error;
	idx_t result_count;
	if (!TryFetchChunk(scan_state, std::move(options), chunk_size, out, result_count, error, extension_type_cast)) {
		error.Throw();
	}
	return result_count;
}

} // namespace duckdb
