/*
 * 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.
 */
#include "paimon/core/io/field_mapping_reader.h"

#include <cassert>
#include <cstddef>
#include <set>
#include <utility>

#include "arrow/api.h"
#include "arrow/array/array_base.h"
#include "arrow/array/array_nested.h"
#include "arrow/array/util.h"
#include "arrow/c/abi.h"
#include "arrow/c/bridge.h"
#include "arrow/scalar.h"
#include "arrow/util/checked_cast.h"
#include "fmt/format.h"
#include "paimon/common/data/binary_string.h"
#include "paimon/common/types/data_field.h"
#include "paimon/common/utils/arrow/mem_utils.h"
#include "paimon/common/utils/arrow/status_utils.h"
#include "paimon/core/casting/cast_executor.h"
#include "paimon/core/casting/casting_utils.h"
#include "paimon/core/utils/field_mapping.h"
#include "paimon/core/utils/nested_projection_utils.h"
#include "paimon/memory/bytes.h"
#include "paimon/reader/batch_reader.h"

namespace paimon {
class MemoryPool;

Result<bool> FieldMappingReader::HasMapSelectedKeysRecursively(
    const std::shared_ptr<arrow::Field>& read_field) const {
    if (!read_field) {
        return false;
    }
    auto type_id = read_field->type()->id();
    if (type_id == arrow::Type::MAP) {
        PAIMON_ASSIGN_OR_RAISE(std::vector<std::string> selected_keys,
                               NestedProjectionUtils::GetMapSelectedKeys(read_field));
        return !selected_keys.empty();
    }
    if (type_id == arrow::Type::STRUCT) {
        for (const auto& child : read_field->type()->fields()) {
            PAIMON_ASSIGN_OR_RAISE(bool has_selected_keys, HasMapSelectedKeysRecursively(child));
            if (has_selected_keys) {
                return true;
            }
        }
    }
    return false;
}

Result<std::shared_ptr<arrow::Array>> FieldMappingReader::FilterMapSelectedKeysRecursively(
    const std::shared_ptr<arrow::Array>& array,
    const std::shared_ptr<arrow::Field>& read_field) const {
    if (!array || !read_field) {
        return array;
    }

    auto type_id = read_field->type()->id();
    if (type_id == arrow::Type::MAP) {
        PAIMON_ASSIGN_OR_RAISE(std::vector<std::string> selected_keys,
                               NestedProjectionUtils::GetMapSelectedKeys(read_field));
        if (selected_keys.empty()) {
            return array;
        }
        return NestedProjectionUtils::FilterMapArrayBySelectedKeys(array, selected_keys,
                                                                   arrow_pool_.get());
    }

    if (type_id == arrow::Type::STRUCT) {
        if (array->type_id() != arrow::Type::STRUCT) {
            return Status::Invalid(
                fmt::format("FilterMapSelectedKeysRecursively requires struct array for read "
                            "field '{}', got {}",
                            read_field->name(), array->type()->ToString()));
        }
        auto struct_array = std::static_pointer_cast<arrow::StructArray>(array);
        auto read_struct_type = std::static_pointer_cast<arrow::StructType>(read_field->type());
        if (struct_array->num_fields() != read_struct_type->num_fields()) {
            return Status::Invalid(fmt::format(
                "FilterMapSelectedKeysRecursively struct field count mismatch for '{}': "
                "array {} vs read {}",
                read_field->name(), struct_array->num_fields(), read_struct_type->num_fields()));
        }

        arrow::ArrayVector filtered_children;
        std::vector<std::shared_ptr<arrow::ArrayData>> filtered_child_data;
        filtered_children.reserve(struct_array->num_fields());
        filtered_child_data.reserve(struct_array->num_fields());
        for (int32_t i = 0; i < struct_array->num_fields(); ++i) {
            PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> filtered_child,
                                   FilterMapSelectedKeysRecursively(struct_array->field(i),
                                                                    read_struct_type->field(i)));
            filtered_child_data.push_back(filtered_child->data());
            filtered_children.push_back(std::move(filtered_child));
        }

        // Preserve parent struct null semantics after filtering children.
        auto filtered_struct_data = arrow::ArrayData::Make(
            read_struct_type, struct_array->length(), {struct_array->null_bitmap()},
            std::move(filtered_child_data), struct_array->null_count(), struct_array->offset());
        return arrow::MakeArray(std::move(filtered_struct_data));
    }

    return array;
}

Result<std::unique_ptr<FieldMappingReader>> FieldMappingReader::Create(
    int32_t field_count, std::unique_ptr<FileBatchReader>&& reader, const BinaryRow& partition,
    std::unique_ptr<FieldMapping>&& mapping,
    std::set<int32_t>&& skip_map_selected_keys_filter_field_ids,
    const std::shared_ptr<MemoryPool>& pool) {
    auto mapping_reader = std::unique_ptr<FieldMappingReader>(new FieldMappingReader(
        field_count, std::move(reader), partition, std::move(mapping), pool));

    mapping_reader->need_mapping_ = false;
    mapping_reader->need_casting_ = false;
    mapping_reader->skip_map_selected_keys_filter_field_ids_ =
        std::move(skip_map_selected_keys_filter_field_ids);

    if (mapping_reader->non_exist_field_info_ != std::nullopt ||
        mapping_reader->partition_info_ != std::nullopt) {
        mapping_reader->need_mapping_ = true;
    }

    for (int32_t i = 0;
         i <
         static_cast<int32_t>(mapping_reader->non_partition_info_.idx_in_target_read_schema.size());
         i++) {
        if (i != mapping_reader->non_partition_info_.idx_in_target_read_schema[i]) {
            mapping_reader->need_mapping_ = true;
        }
        if (mapping_reader->non_partition_info_.cast_executors[i] != nullptr) {
            mapping_reader->need_casting_ = true;
        }
        // A differing nested type needs the AlignArrayToReadType reshape below.
        if (!mapping_reader->non_partition_info_.non_partition_data_schema[i].Type()->Equals(
                *mapping_reader->non_partition_info_.non_partition_read_schema[i].Type())) {
            mapping_reader->need_casting_ = true;
        }
        // Field name change (RENAME COLUMN) also requires mapping: data schema
        // carries the file's physical name while read schema carries the
        // post-rename logical name. If we skipped mapping, the inner reader's
        // batch would be passed through with the old physical name and the
        // consumer's name-based lookup against the read schema would fail.
        if (mapping_reader->non_partition_info_.non_partition_data_schema[i].Name() !=
            mapping_reader->non_partition_info_.non_partition_read_schema[i].Name()) {
            mapping_reader->need_mapping_ = true;
        }
        // Map selected-keys metadata must be validated in Create() (fail-fast).
        // Non-empty selected-keys also requires mapping so that
        // FilterMapArrayBySelectedKeys can filter out unwanted entries.
        PAIMON_ASSIGN_OR_RAISE(
            bool has_map_selected_keys,
            mapping_reader->HasMapSelectedKeysRecursively(
                mapping_reader->non_partition_info_.non_partition_read_schema[i].ArrowField()));
        if (has_map_selected_keys &&
            mapping_reader->skip_map_selected_keys_filter_field_ids_.count(
                mapping_reader->non_partition_info_.non_partition_read_schema[i].Id()) == 0) {
            mapping_reader->need_mapping_ = true;
        }
    }

    return mapping_reader;
}

FieldMappingReader::FieldMappingReader(int32_t field_count,
                                       std::unique_ptr<FileBatchReader>&& reader,
                                       const BinaryRow& partition,
                                       std::unique_ptr<FieldMapping>&& mapping,
                                       const std::shared_ptr<MemoryPool>& pool)
    : field_count_(field_count),
      arrow_pool_(GetArrowPool(pool)),
      reader_(std::move(reader)),
      partition_(partition),
      partition_info_(mapping->partition_info),
      non_partition_info_(mapping->non_partition_info),
      non_exist_field_info_(mapping->non_exist_field_info) {}

Result<std::shared_ptr<arrow::Array>> FieldMappingReader::CastNonPartitionArrayIfNeed(
    const std::shared_ptr<arrow::Array>& src_array) const {
    if (!need_casting_) {
        return src_array;
    }
    auto* struct_array = arrow::internal::checked_cast<arrow::StructArray*>(src_array.get());
    int32_t field_count = struct_array->num_fields();
    assert(static_cast<size_t>(field_count) == non_partition_info_.cast_executors.size());
    arrow::ArrayVector casted_array;
    std::vector<std::string> casted_field_names;
    casted_array.reserve(field_count);
    casted_field_names.reserve(field_count);
    for (int32_t i = 0; i < field_count; i++) {
        std::shared_ptr<arrow::Array> column;
        if (non_partition_info_.cast_executors[i] != nullptr) {
            auto single_column_array = struct_array->field(i);
            // if src array is dict, cast to string first
            auto dict_array =
                std::dynamic_pointer_cast<arrow::DictionaryArray>(single_column_array);
            if (dict_array) {
                PAIMON_ASSIGN_OR_RAISE(
                    single_column_array,
                    CastingUtils::Cast(dict_array, /*target_type=*/arrow::utf8(),
                                       arrow::compute::CastOptions::Safe(), arrow_pool_.get()));
            }
            PAIMON_ASSIGN_OR_RAISE(
                column,
                non_partition_info_.cast_executors[i]->Cast(
                    single_column_array, non_partition_info_.non_partition_read_schema[i].Type(),
                    arrow_pool_.get()));
        } else {
            // read and data type may both be string type, but after adapter transform, type may be
            // dictionary, need reconstruct struct type
            column = struct_array->field(i);
        }
        // Null-fill nested fields added by schema evolution. Only when the data and
        // read types differ -- the reader may hand back a dictionary-encoded array
        // for an unchanged type, which is not a reshape target.
        if (!non_partition_info_.non_partition_data_schema[i].Type()->Equals(
                *non_partition_info_.non_partition_read_schema[i].Type())) {
            PAIMON_ASSIGN_OR_RAISE(
                column, NestedProjectionUtils::AlignArrayToReadType(
                            column, non_partition_info_.non_partition_read_schema[i].Type(),
                            arrow_pool_.get()));
        }
        casted_array.push_back(column);
        casted_field_names.push_back(non_partition_info_.non_partition_data_schema[i].Name());
    }
    PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> arrow_array,
                                      arrow::StructArray::Make(casted_array, casted_field_names));
    return arrow_array;
}

Result<BatchReader::ReadBatchWithBitmap> FieldMappingReader::NextBatchWithBitmap() {
    PAIMON_ASSIGN_OR_RAISE(ReadBatchWithBitmap non_partition_result_with_bitmap,
                           reader_->NextBatchWithBitmap());
    if (!need_mapping_ && !need_casting_) {
        return non_partition_result_with_bitmap;
    }
    if (BatchReader::IsEofBatch(non_partition_result_with_bitmap)) {
        // read finish
        partition_array_.reset();
        non_exist_array_.reset();
        return non_partition_result_with_bitmap;
    }
    auto& [non_partition_result, bitmap] = non_partition_result_with_bitmap;
    auto& [c_array, c_schema] = non_partition_result;
    PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> non_partition_array,
                                      arrow::ImportArray(c_array.get(), c_schema.get()));

    arrow::ArrayVector target_array(field_count_);
    std::vector<std::string> target_field_names(field_count_);
    // mapping non-partition array
    PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> casted_non_partition_array,
                           CastNonPartitionArrayIfNeed(non_partition_array));
    PAIMON_RETURN_NOT_OK(MappingFields(
        casted_non_partition_array, non_partition_info_.non_partition_read_schema,
        non_partition_info_.idx_in_target_read_schema, &target_array, &target_field_names));

    // mapping partition array
    if (partition_info_ != std::nullopt) {
        if (!partition_array_ || partition_array_->length() < non_partition_array->length()) {
            PAIMON_ASSIGN_OR_RAISE(partition_array_,
                                   GeneratePartitionArray(non_partition_array->length()));
        }
        auto trim_partition_array = partition_array_->Slice(0, non_partition_array->length());
        PAIMON_RETURN_NOT_OK(MappingFields(
            trim_partition_array, partition_info_.value().partition_read_schema,
            partition_info_.value().idx_in_target_read_schema, &target_array, &target_field_names));
    }
    // mapping non-exist array
    if (non_exist_field_info_ != std::nullopt) {
        if (!non_exist_array_ || non_exist_array_->length() < non_partition_array->length()) {
            PAIMON_ASSIGN_OR_RAISE(non_exist_array_,
                                   GenerateNonExistArray(non_partition_array->length()));
        }
        auto trim_non_exist_array = non_exist_array_->Slice(0, non_partition_array->length());
        PAIMON_RETURN_NOT_OK(MappingFields(trim_non_exist_array,
                                           non_exist_field_info_.value().non_exist_read_schema,
                                           non_exist_field_info_.value().idx_in_target_read_schema,
                                           &target_array, &target_field_names));
    }

    // construct target array
    PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> arrow_array,
                                      arrow::StructArray::Make(target_array, target_field_names));
    std::unique_ptr<ArrowArray> target_c_arrow_array = std::make_unique<ArrowArray>();
    std::unique_ptr<ArrowSchema> target_c_schema = std::make_unique<ArrowSchema>();
    PAIMON_RETURN_NOT_OK_FROM_ARROW(
        arrow::ExportArray(*arrow_array, target_c_arrow_array.get(), target_c_schema.get()));
    auto target_batch = std::make_pair(std::move(target_c_arrow_array), std::move(target_c_schema));
    return std::make_pair(std::move(target_batch), std::move(bitmap));
}

Result<std::shared_ptr<arrow::Array>> FieldMappingReader::GenerateSinglePartitionArray(
    int32_t idx, int32_t batch_size) const {
    const auto& type = partition_info_.value().partition_read_schema[idx].Type();
    if (partition_.IsNullAt(partition_info_.value().idx_in_partition[idx])) {
        // for null partition value
        PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
            std::shared_ptr<arrow::Array> null_array,
            arrow::MakeArrayOfNull(type, batch_size, arrow_pool_.get()));
        return null_array;
    }
    auto type_id = type->id();
    std::shared_ptr<arrow::Scalar> scalar;
    switch (type_id) {
        case arrow::Type::type::BOOL: {
            bool value = partition_.GetBoolean(partition_info_.value().idx_in_partition[idx]);
            scalar = std::make_shared<arrow::BooleanScalar>(value);
            break;
        }
        case arrow::Type::type::INT8: {
            int8_t value = partition_.GetByte(partition_info_.value().idx_in_partition[idx]);
            scalar = std::make_shared<arrow::Int8Scalar>(value);
            break;
        }
        case arrow::Type::type::INT16: {
            int16_t value = partition_.GetShort(partition_info_.value().idx_in_partition[idx]);
            scalar = std::make_shared<arrow::Int16Scalar>(value);
            break;
        }
        case arrow::Type::type::INT32: {
            int32_t value = partition_.GetInt(partition_info_.value().idx_in_partition[idx]);
            scalar = std::make_shared<arrow::Int32Scalar>(value);
            break;
        }
        case arrow::Type::type::INT64: {
            int64_t value = partition_.GetLong(partition_info_.value().idx_in_partition[idx]);
            scalar = std::make_shared<arrow::Int64Scalar>(value);
            break;
        }
        case arrow::Type::type::STRING: {
            BinaryString value =
                partition_.GetString(partition_info_.value().idx_in_partition[idx]);
            scalar = std::make_shared<arrow::StringScalar>(value.ToString());
            break;
        }
        case arrow::Type::type::BINARY: {
            auto value = partition_.GetBinary(partition_info_.value().idx_in_partition[idx]);
            std::string value_str(value->data(), value->size());
            scalar = std::make_shared<arrow::BinaryScalar>(value_str);
            break;
        }
        case arrow::Type::type::DATE32: {
            int32_t value = partition_.GetDate(partition_info_.value().idx_in_partition[idx]);
            scalar = std::make_shared<arrow::Date32Scalar>(value);
            break;
        }
        default:
            return Status::Invalid(
                fmt::format("Not support arrow type {} for partition",
                            partition_info_.value().partition_read_schema[idx].Type()->ToString()));
    }
    PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
        std::shared_ptr<arrow::Array> arrow_array,
        arrow::MakeArrayFromScalar(*scalar, batch_size, arrow_pool_.get()));
    return arrow_array;
}

Result<std::shared_ptr<arrow::Array>> FieldMappingReader::GeneratePartitionArray(
    int32_t batch_size) const {
    arrow::ArrayVector partition_array;
    std::vector<std::string> partition_field_names;
    partition_array.reserve(partition_info_.value().partition_read_schema.size());
    partition_field_names.reserve(partition_info_.value().partition_read_schema.size());
    for (size_t i = 0; i < partition_info_.value().partition_read_schema.size(); i++) {
        PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> single_partition_array,
                               GenerateSinglePartitionArray(i, batch_size));
        partition_array.push_back(single_partition_array);
        partition_field_names.push_back(partition_info_.value().partition_read_schema[i].Name());
    }
    PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
        std::shared_ptr<arrow::Array> arrow_array,
        arrow::StructArray::Make(partition_array, partition_field_names));
    return arrow_array;
}

Result<std::shared_ptr<arrow::Array>> FieldMappingReader::GenerateNonExistArray(
    int32_t batch_size) const {
    arrow::ArrayVector non_exist_array;
    std::vector<std::string> non_exist_field_names;
    non_exist_array.reserve(non_exist_field_info_.value().non_exist_read_schema.size());
    non_exist_field_names.reserve(non_exist_field_info_.value().non_exist_read_schema.size());
    for (const auto& non_exist_field : non_exist_field_info_.value().non_exist_read_schema) {
        PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
            std::shared_ptr<arrow::Array> null_array,
            arrow::MakeArrayOfNull(non_exist_field.Type(), batch_size, arrow_pool_.get()));
        non_exist_array.push_back(null_array);
        non_exist_field_names.push_back(non_exist_field.Name());
    }
    PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
        std::shared_ptr<arrow::Array> arrow_array,
        arrow::StructArray::Make(non_exist_array, non_exist_field_names));
    return arrow_array;
}

Status FieldMappingReader::MappingFields(const std::shared_ptr<arrow::Array>& data_array,
                                         const std::vector<DataField>& read_fields_of_data_array,
                                         const std::vector<int32_t>& idx_in_target_schema,
                                         arrow::ArrayVector* target_array,
                                         std::vector<std::string>* target_field_names) {
    auto* struct_array = arrow::internal::checked_cast<arrow::StructArray*>(data_array.get());
    assert(struct_array);
    assert(struct_array->fields().size() == idx_in_target_schema.size());
    for (size_t i = 0; i < idx_in_target_schema.size(); i++) {
        std::shared_ptr<arrow::Array> field_array = struct_array->field(i);
        const DataField& read_field = read_fields_of_data_array[i];

        // Filter map entries by selected keys recursively (supports MAP nested in STRUCT).
        if (skip_map_selected_keys_filter_field_ids_.count(read_field.Id()) == 0) {
            PAIMON_ASSIGN_OR_RAISE(field_array, FilterMapSelectedKeysRecursively(
                                                    field_array, read_field.ArrowField()));
        }

        (*target_array)[idx_in_target_schema[i]] = std::move(field_array);
        (*target_field_names)[idx_in_target_schema[i]] = read_field.Name();
    }
    return Status::OK();
}

}  // namespace paimon
