/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

#include <unistd.h>

#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "arrow/api.h"
#include "arrow/c/bridge.h"
#include "arrow/compute/api.h"
#include "paimon/common/reader/reader_utils.h"
#include "paimon/common/utils/arrow/status_utils.h"
#include "paimon/common/utils/date_time_utils.h"
#include "paimon/core/io/key_value_data_file_record_reader.h"
#include "paimon/core/key_value.h"
#include "paimon/reader/batch_reader.h"
#include "paimon/testing/utils/dict_array_converter.h"
namespace paimon::test {
class ReadResultCollector {
 public:
    ReadResultCollector() = delete;
    ~ReadResultCollector() = delete;

    template <typename ReaderType, typename IteratorType>
    static Result<std::vector<KeyValue>> CollectKeyValueResult(ReaderType* reader) {
        std::vector<KeyValue> results;
        while (true) {
            PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<IteratorType> iterator, reader->NextBatch());
            if (iterator == nullptr) {
                break;
            }
            while (true) {
                if constexpr (std::is_same_v<IteratorType, KeyValueRecordReader::Iterator>) {
                    PAIMON_ASSIGN_OR_RAISE(bool has_next, iterator->HasNext());
                    if (!has_next) {
                        break;
                    }
                    PAIMON_ASSIGN_OR_RAISE(KeyValue kv, iterator->Next());
                    results.emplace_back(std::move(kv));
                } else {
                    PAIMON_ASSIGN_OR_RAISE(bool has_next, iterator->HasNext());
                    if (!has_next) {
                        break;
                    }
                    KeyValue kv = iterator->Next();
                    results.emplace_back(std::move(kv));
                }
            }
        }
        return results;
    }

    // will convert dictionary array to string array for comparing results
    static Result<std::shared_ptr<arrow::ChunkedArray>> CollectResult(
        BatchReader* batch_reader, int64_t max_data_processing_time_in_us) {
        arrow::ArrayVector result_array_vector;
        int64_t seed = DateTimeUtils::GetCurrentUTCTimeUs();
        std::srand(seed);
        while (true) {
            PAIMON_ASSIGN_OR_RAISE(auto result_array, ReadOneBatch(batch_reader));
            if (result_array == nullptr) {
                break;
            }
            result_array_vector.push_back(result_array);
            if (max_data_processing_time_in_us > 0) {
                usleep(std::rand() % max_data_processing_time_in_us);
            }
        }
        if (result_array_vector.empty()) {
            return std::shared_ptr<arrow::ChunkedArray>();
        }
        // accumulate all the batch array and convert dictionary to string array together to avoid
        // the problem (multiple batches in multiple stripes overlap dictionary data) being
        // difficult to expose
        arrow::ArrayVector converted_array_vector;
        for (const auto& array : result_array_vector) {
            PAIMON_ASSIGN_OR_RAISE(
                std::shared_ptr<arrow::Array> converted_array,
                DictArrayConverter::ConvertDictArray(array, arrow::default_memory_pool()));
            converted_array_vector.push_back(converted_array);
        }
        PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(auto chunk_array,
                                          arrow::ChunkedArray::Make(converted_array_vector));
        return chunk_array;
    }

    static Result<std::shared_ptr<arrow::ChunkedArray>> CollectResult(BatchReader* batch_reader) {
        return CollectResult(batch_reader, /*max_data_processing_time_in_us=*/0);
    }

    static Result<std::shared_ptr<arrow::ChunkedArray>> CollectResultOneBatch(
        BatchReader* batch_reader) {
        return CollectResultOneBatch(batch_reader, /*max_data_processing_time_in_us=*/0);
    }

    static Result<std::shared_ptr<arrow::ChunkedArray>> CollectResultOneBatch(
        BatchReader* batch_reader, int64_t max_data_processing_time_in_us) {
        int64_t seed = DateTimeUtils::GetCurrentUTCTimeUs();
        std::srand(seed);
        PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> result_array,
                               ReadOneBatch(batch_reader));
        if (result_array == nullptr) {
            return std::shared_ptr<arrow::ChunkedArray>();
        }
        PAIMON_ASSIGN_OR_RAISE(
            auto converted_array,
            DictArrayConverter::ConvertDictArray(result_array, arrow::default_memory_pool()));
        if (max_data_processing_time_in_us > 0) {
            usleep(std::rand() % max_data_processing_time_in_us);
        }
        PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::ChunkedArray> chunk_array,
                                          arrow::ChunkedArray::Make({converted_array}));
        return chunk_array;
    }

    static Result<std::shared_ptr<arrow::Array>> GetArray(BatchReader::ReadBatch&& batch) {
        if (BatchReader::IsEofBatch(batch)) {
            return std::shared_ptr<arrow::Array>();
        }
        auto& [c_array, c_schema] = batch;
        PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(auto array,
                                          arrow::ImportArray(c_array.get(), c_schema.get()));
        return DictArrayConverter::ConvertDictArray(array, arrow::default_memory_pool());
    }

    static Result<BatchReader::ReadBatch> GetReadBatch(const std::shared_ptr<arrow::Array>& array) {
        auto c_array = std::make_unique<ArrowArray>();
        auto c_schema = std::make_unique<ArrowSchema>();
        PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, c_array.get(), c_schema.get()));
        return std::make_pair(std::move(c_array), std::move(c_schema));
    }

    // Noted that, sort chunked array by multiple key for timestamp type may cause
    // coredump in arrow, refer to https://github.com/apache/arrow/issues/47252
    static Result<std::shared_ptr<arrow::ChunkedArray>> SortArray(
        const std::shared_ptr<arrow::ChunkedArray>& array,
        const std::shared_ptr<arrow::Schema>& schema) {
        std::vector<arrow::compute::SortKey> sort_keys;
        for (const auto& name : schema->field_names()) {
            sort_keys.emplace_back(name, arrow::compute::SortOrder::Ascending);
        }
        auto sort_options = arrow::compute::SortOptions(sort_keys);
        PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
            auto sorted_indices, arrow::compute::SortIndices(arrow::Datum(array), sort_options));
        PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
            auto sorted_batch,
            arrow::compute::Take(arrow::Datum(array), arrow::Datum(sorted_indices)));
        return sorted_batch.chunked_array();
    }

 private:
    static Result<std::shared_ptr<arrow::Array>> ReadOneBatch(BatchReader* batch_reader) {
        // Prioritize calling NextBatch. If it fails (paimon inner reader e.g.,
        // PrefetchBatchReader, ApplyBitmapIndexBatchReader...), call NextBatchWithBitmap.
        auto batch_result = batch_reader->NextBatch();
        BatchReader::ReadBatch batch;
        if (!batch_result.ok()) {
            if (batch_result.status().ToString().find("should use NextBatchWithBitmap") !=
                std::string::npos) {
                PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatchWithBitmap batch_with_bitmap,
                                       batch_reader->NextBatchWithBitmap());
                if (BatchReader::IsEofBatch(batch_with_bitmap)) {
                    return std::shared_ptr<arrow::Array>();
                }
                PAIMON_RETURN_NOT_OK(CheckBatchOffset(batch_with_bitmap.first));
                assert(!batch_with_bitmap.second.IsEmpty());
                PAIMON_ASSIGN_OR_RAISE(
                    batch, ReaderUtils::ApplyBitmapToReadBatch(std::move(batch_with_bitmap),
                                                               arrow::default_memory_pool()));
            } else {
                return batch_result.status();
            }
        } else {
            batch = std::move(batch_result).value();
            if (BatchReader::IsEofBatch(batch)) {
                return std::shared_ptr<arrow::Array>();
            }
            PAIMON_RETURN_NOT_OK(CheckBatchOffset(batch));
        }
        auto& [c_array, c_schema] = batch;
        assert(c_array->length > 0);
        PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(auto result_array,
                                          arrow::ImportArray(c_array.get(), c_schema.get()));
        return result_array;
    }

    static Status CheckBatchOffset(const BatchReader::ReadBatch& batch) {
        assert(!BatchReader::IsEofBatch(batch));
        return CheckArrayOffset(batch.first.get());
    }

    static Status CheckArrayOffset(const ArrowArray* array) {
        assert(array);
        if (array->offset != 0) {
            return Status::Invalid("BatchReader returned an array with non-zero offset " +
                                   std::to_string(array->offset));
        }
        for (int64_t i = 0; i < array->n_children; i++) {
            PAIMON_RETURN_NOT_OK(CheckArrayOffset(array->children[i]));
        }
        return Status::OK();
    }
};
}  // namespace paimon::test
