/*
 * 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/common/data/shredding/map_shared_shredding_utils.h"

#include <algorithm>
#include <functional>

#include "arrow/type.h"
#include "arrow/util/key_value_metadata.h"
#include "fmt/format.h"
#include "paimon/common/compression/block_compression_factory.h"
#include "paimon/common/compression/block_compressor.h"
#include "paimon/common/compression/block_decompressor.h"
#include "paimon/common/data/shredding/map_shared_shredding_batch_converter.h"
#include "paimon/common/utils/arrow/status_utils.h"
#include "paimon/common/utils/string_utils.h"
#include "paimon/core/core_options.h"
#include "paimon/core/options/map_storage_layout.h"
#include "rapidjson/document.h"
#include "rapidjson/stringbuffer.h"
#include "rapidjson/writer.h"

namespace paimon {
Result<std::vector<int32_t>> MapSharedShreddingUtils::GetPhysicalColumnIndices(
    const MapSharedShreddingFieldMeta& meta, const std::string& name) {
    auto name_iter = meta.name_to_id.find(name);
    if (name_iter == meta.name_to_id.end()) {
        return Status::Invalid(
            fmt::format("cannot find field {} in map shared shredding meta", name));
    }
    auto id_iter = meta.field_to_columns.find(name_iter->second);
    if (id_iter == meta.field_to_columns.end()) {
        return Status::Invalid(
            fmt::format("cannot find field id {} in field_to_columns in map shared shredding meta",
                        name_iter->second));
    }
    return id_iter->second;
}

// ---- Column detection ----
bool MapSharedShreddingUtils::IsShreddingKeyMap(
    const std::shared_ptr<arrow::DataType>& arrow_type) {
    if (arrow_type->id() != arrow::Type::MAP) {
        return false;
    }
    auto map_type = std::static_pointer_cast<arrow::MapType>(arrow_type);
    return map_type->key_type()->id() == arrow::Type::STRING;
}

Result<std::vector<std::string>> MapSharedShreddingUtils::DetectShreddingColumns(
    const std::shared_ptr<arrow::Schema>& schema, const CoreOptions& options) {
    std::vector<std::string> field_names;
    for (int32_t i = 0; i < schema->num_fields(); ++i) {
        const auto& field = schema->field(i);
        if (!IsShreddingKeyMap(field->type())) {
            continue;
        }
        PAIMON_ASSIGN_OR_RAISE(MapStorageLayout layout, options.GetMapStorageLayout(field->name()));
        if (layout == MapStorageLayout::SHARED_SHREDDING) {
            field_names.push_back(field->name());
        }
    }
    return field_names;
}

// ---- Schema conversion ----
std::shared_ptr<arrow::DataType> MapSharedShreddingUtils::BuildSpecificPhysicalStructType(
    const std::shared_ptr<arrow::DataType>& value_type, const std::set<int32_t>& physical_col_ids,
    bool value_nullable, bool include_overflow) {
    std::vector<int32_t> sorted_cols(physical_col_ids.begin(), physical_col_ids.end());
    return InnerBuildSpecificPhysicalStructType(value_type, sorted_cols, value_nullable,
                                                include_overflow);
}

std::shared_ptr<arrow::DataType> MapSharedShreddingUtils::BuildPhysicalStructType(
    const std::shared_ptr<arrow::DataType>& value_type, int32_t num_columns, bool value_nullable) {
    std::vector<int32_t> sorted_cols(num_columns);
    std::iota(sorted_cols.begin(), sorted_cols.end(), 0);
    return InnerBuildSpecificPhysicalStructType(value_type, sorted_cols, value_nullable,
                                                /*include_overflow=*/true);
}

std::shared_ptr<arrow::DataType> MapSharedShreddingUtils::InnerBuildSpecificPhysicalStructType(
    const std::shared_ptr<arrow::DataType>& value_type, const std::vector<int32_t>& sorted_cols,
    bool value_nullable, bool include_overflow) {
    arrow::FieldVector struct_fields;
    struct_fields.reserve(sorted_cols.size() + 2);
    struct_fields.push_back(
        arrow::field(MapSharedShreddingDefine::kFieldMapping, arrow::list(arrow::int32()), true));
    for (const auto& col : sorted_cols) {
        struct_fields.push_back(arrow::field(MapSharedShreddingDefine::PhysicalColumnName(col),
                                             value_type, value_nullable));
    }
    if (include_overflow) {
        struct_fields.push_back(arrow::field(
            MapSharedShreddingDefine::kOverflow,
            arrow::map(arrow::int32(), arrow::field("value", value_type, value_nullable)), true));
    }
    return arrow::struct_(std::move(struct_fields));
}

Result<std::shared_ptr<arrow::Schema>> MapSharedShreddingUtils::LogicalToPhysicalSchema(
    const std::shared_ptr<arrow::Schema>& logical_schema,
    const std::map<std::string, int32_t>& field_to_num_columns) {
    arrow::FieldVector physical_fields;
    physical_fields.reserve(logical_schema->num_fields());

    for (int32_t i = 0; i < logical_schema->num_fields(); ++i) {
        const auto& field = logical_schema->field(i);
        auto it = field_to_num_columns.find(field->name());
        if (it != field_to_num_columns.end()) {
            if (field->type()->id() != arrow::Type::MAP) {
                return Status::Invalid(
                    fmt::format("Field '{}' is expected to be MAP type, but got '{}'.",
                                field->name(), field->type()->name()));
            }
            auto map_type = std::static_pointer_cast<arrow::MapType>(field->type());
            auto value_type = map_type->item_type();
            bool value_nullable = map_type->item_field()->nullable();
            auto physical_type = BuildPhysicalStructType(value_type, it->second, value_nullable);
            auto physical_field = field->WithType(physical_type);
            physical_fields.push_back(physical_field);
        } else {
            physical_fields.push_back(field);
        }
    }

    return arrow::schema(std::move(physical_fields));
}

Result<std::map<std::string, int32_t>> MapSharedShreddingUtils::BuildColumnToNumColumns(
    const std::vector<std::string>& shredding_field_names, const CoreOptions& options) {
    std::map<std::string, int32_t> field_to_num_columns;
    for (const std::string& field_name : shredding_field_names) {
        PAIMON_ASSIGN_OR_RAISE(int32_t max_columns,
                               options.GetMapSharedShreddingMaxColumns(field_name));
        field_to_num_columns[field_name] = max_columns;
    }
    return field_to_num_columns;
}

// ---- Metadata serialization helpers ----

namespace {

std::string JsonEncodeObject(
    std::function<void(rapidjson::Document*, rapidjson::Document::AllocatorType*)> builder) {
    rapidjson::Document doc(rapidjson::kObjectType);
    auto& allocator = doc.GetAllocator();
    builder(&doc, &allocator);
    rapidjson::StringBuffer buffer;
    rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);
    doc.Accept(writer);
    return buffer.GetString();
}

std::string JsonEncodeArray(
    std::function<void(rapidjson::Document*, rapidjson::Document::AllocatorType*)> builder) {
    rapidjson::Document doc(rapidjson::kArrayType);
    auto& allocator = doc.GetAllocator();
    builder(&doc, &allocator);
    rapidjson::StringBuffer buffer;
    rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);
    doc.Accept(writer);
    return buffer.GetString();
}

Result<std::string> CompressString(const std::string& input, const std::string& compression) {
    CompressOptions compress_opts{compression, /*zstd_level=*/1};
    PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<BlockCompressionFactory> factory,
                           BlockCompressionFactory::Create(compress_opts));
    std::shared_ptr<BlockCompressor> compressor = factory->GetCompressor();
    if (!compressor) {
        return input;
    }

    auto src_size = static_cast<int32_t>(input.size());
    int32_t max_compressed = compressor->GetMaxCompressedSize(src_size);
    std::string output(max_compressed, '\0');

    PAIMON_ASSIGN_OR_RAISE(
        int32_t actual_size,
        compressor->Compress(input.data(), src_size, output.data(), max_compressed));

    output.resize(actual_size);
    return output;
}

Result<std::string> DecompressString(const std::string& input, int32_t original_len,
                                     const std::string& compression) {
    CompressOptions compress_opts{compression, /*zstd_level=*/1};
    PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<BlockCompressionFactory> factory,
                           BlockCompressionFactory::Create(compress_opts));
    std::shared_ptr<BlockDecompressor> decompressor = factory->GetDecompressor();
    if (!decompressor) {
        return input;
    }
    std::string output(original_len, '\0');
    PAIMON_ASSIGN_OR_RAISE(
        int32_t decompressed_len,
        decompressor->Decompress(input.data(), static_cast<int32_t>(input.size()), output.data(),
                                 original_len));
    output.resize(decompressed_len);
    return output;
}

Result<std::string> GetRequiredValue(const std::shared_ptr<arrow::KeyValueMetadata>& metadata,
                                     const std::string& key) {
    int32_t index = metadata->FindKey(key);
    if (index < 0) {
        return Status::Invalid(fmt::format("missing shredding metadata key: {}", key));
    }
    return metadata->value(index);
}

Result<int32_t> GetRequiredInt32(const std::shared_ptr<arrow::KeyValueMetadata>& metadata,
                                 const std::string& key) {
    PAIMON_ASSIGN_OR_RAISE(std::string value, GetRequiredValue(metadata, key));
    std::optional<int32_t> parsed = StringUtils::StringToValue<int32_t>(value);
    if (!parsed.has_value()) {
        return Status::Invalid(fmt::format("malformed shredding metadata value for key: {}", key));
    }
    return parsed.value();
}

std::string SerializeFieldDict(const MapSharedShreddingFieldMeta& field_meta) {
    return JsonEncodeObject([&](rapidjson::Document* doc,
                                rapidjson::Document::AllocatorType* alloc) {
        for (const auto& [name, id] : field_meta.name_to_id) {
            doc->AddMember(rapidjson::Value(name.c_str(), *alloc), rapidjson::Value(id), *alloc);
        }
    });
}

std::string SerializeFieldColumns(const MapSharedShreddingFieldMeta& field_meta) {
    return JsonEncodeObject(
        [&](rapidjson::Document* doc, rapidjson::Document::AllocatorType* alloc) {
            for (const auto& [field_id, col_vec] : field_meta.field_to_columns) {
                rapidjson::Value array(rapidjson::kArrayType);
                std::vector<int32_t> sorted_cols(col_vec.begin(), col_vec.end());
                std::sort(sorted_cols.begin(), sorted_cols.end());
                for (int32_t col : sorted_cols) {
                    array.PushBack(col, *alloc);
                }
                std::string key = std::to_string(field_id);
                doc->AddMember(rapidjson::Value(key.c_str(), *alloc), array, *alloc);
            }
        });
}

std::string SerializeOverflowSet(const MapSharedShreddingFieldMeta& field_meta) {
    return JsonEncodeArray(
        [&](rapidjson::Document* doc, rapidjson::Document::AllocatorType* alloc) {
            std::vector<int32_t> sorted(field_meta.overflow_field_set.begin(),
                                        field_meta.overflow_field_set.end());
            std::sort(sorted.begin(), sorted.end());
            for (int32_t field_id : sorted) {
                doc->PushBack(field_id, *alloc);
            }
        });
}

/// Safe JSON integer extraction with error propagation.
Result<int32_t> JsonGetInt(const rapidjson::Value& val, const char* context_msg) {
    if (!val.IsInt()) {
        return Status::Invalid(fmt::format("malformed shredding metadata: {}", context_msg));
    }
    return val.GetInt();
}

Result<std::map<std::string, int32_t>> DeserializeFieldDict(const std::string& json_str) {
    rapidjson::Document doc;
    doc.Parse(json_str.c_str());
    if (doc.HasParseError() || !doc.IsObject()) {
        return Status::Invalid("malformed shredding field_dict metadata");
    }
    std::map<std::string, int32_t> name_to_id;
    for (auto it = doc.MemberBegin(); it != doc.MemberEnd(); ++it) {
        PAIMON_ASSIGN_OR_RAISE(int32_t id, JsonGetInt(it->value, "field_dict value is not int"));
        name_to_id[it->name.GetString()] = id;
    }
    return name_to_id;
}

Result<std::map<int32_t, std::vector<int32_t>>> DeserializeFieldColumns(
    const std::string& json_str) {
    rapidjson::Document doc;
    doc.Parse(json_str.c_str());
    if (doc.HasParseError() || !doc.IsObject()) {
        return Status::Invalid("malformed shredding field_columns metadata");
    }
    std::map<int32_t, std::vector<int32_t>> field_to_columns;
    for (auto it = doc.MemberBegin(); it != doc.MemberEnd(); ++it) {
        std::optional<int32_t> field_id = StringUtils::StringToValue<int32_t>(it->name.GetString());
        if (!field_id.has_value()) {
            return Status::Invalid("malformed shredding field_columns: invalid field_id key");
        }
        const auto& array = it->value;
        if (!array.IsArray()) {
            return Status::Invalid("malformed shredding field_columns: value is not array");
        }
        std::vector<int32_t> cols;
        cols.reserve(array.Size());
        for (rapidjson::SizeType i = 0; i < array.Size(); ++i) {
            PAIMON_ASSIGN_OR_RAISE(int32_t col,
                                   JsonGetInt(array[i], "field_columns element is not int"));
            cols.push_back(col);
        }
        field_to_columns[field_id.value()] = std::move(cols);
    }
    return field_to_columns;
}

Result<std::set<int32_t>> DeserializeOverflowSet(const std::string& json_str) {
    rapidjson::Document doc;
    doc.Parse(json_str.c_str());
    if (doc.HasParseError() || !doc.IsArray()) {
        return Status::Invalid("malformed shredding overflow_set metadata");
    }
    std::set<int32_t> overflow_set;
    for (rapidjson::SizeType i = 0; i < doc.Size(); ++i) {
        PAIMON_ASSIGN_OR_RAISE(int32_t field_id,
                               JsonGetInt(doc[i], "overflow_set element is not int"));
        overflow_set.insert(field_id);
    }
    return overflow_set;
}

}  // namespace

Status MapSharedShreddingUtils::SerializeMetadata(const MapSharedShreddingFieldMeta& field_meta,
                                                  const std::string& compression,
                                                  arrow::KeyValueMetadata* metadata) {
    const std::string normalized_compression = StringUtils::ToLowerCase(compression);
    PAIMON_RETURN_NOT_OK_FROM_ARROW(metadata->Set(
        MapShreddingDefine::kStorageLayout, MapShreddingDefine::kStorageLayoutSharedShredding));
    PAIMON_RETURN_NOT_OK_FROM_ARROW(
        metadata->Set(MapSharedShreddingDefine::kVersion,
                      std::to_string(MapSharedShreddingDefine::kCurrentVersion)));

    std::string field_dict_json = SerializeFieldDict(field_meta);
    PAIMON_RETURN_NOT_OK_FROM_ARROW(metadata->Set(MapSharedShreddingDefine::kFieldDictOriginalSize,
                                                  std::to_string(field_dict_json.size())));
    PAIMON_RETURN_NOT_OK_FROM_ARROW(
        metadata->Set(MapSharedShreddingDefine::kFieldDictCompression, normalized_compression));
    PAIMON_ASSIGN_OR_RAISE(std::string compressed_dict,
                           CompressString(field_dict_json, normalized_compression));
    PAIMON_RETURN_NOT_OK_FROM_ARROW(
        metadata->Set(MapSharedShreddingDefine::kFieldDict, std::move(compressed_dict)));

    PAIMON_RETURN_NOT_OK_FROM_ARROW(
        metadata->Set(MapSharedShreddingDefine::kFieldColumns, SerializeFieldColumns(field_meta)));
    PAIMON_RETURN_NOT_OK_FROM_ARROW(
        metadata->Set(MapSharedShreddingDefine::kOverflowSet, SerializeOverflowSet(field_meta)));
    PAIMON_RETURN_NOT_OK_FROM_ARROW(metadata->Set(MapSharedShreddingDefine::kNumColumns,
                                                  std::to_string(field_meta.num_columns)));
    PAIMON_RETURN_NOT_OK_FROM_ARROW(metadata->Set(MapSharedShreddingDefine::kMaxRowWidth,
                                                  std::to_string(field_meta.max_row_width)));

    return Status::OK();
}

Result<MapSharedShreddingFieldMeta> MapSharedShreddingUtils::DeserializeMetadata(
    const std::shared_ptr<arrow::KeyValueMetadata>& metadata) {
    if (!HasShreddingMetadata(metadata)) {
        return Status::Invalid("metadata is null or storage layout is not shared-shredding");
    }
    PAIMON_ASSIGN_OR_RAISE(int32_t version,
                           GetRequiredInt32(metadata, MapSharedShreddingDefine::kVersion));
    if (version != MapSharedShreddingDefine::kCurrentVersion) {
        return Status::Invalid(
            fmt::format("unsupported shared-shredding metadata version: {}, expected: {}", version,
                        MapSharedShreddingDefine::kCurrentVersion));
    }

    MapSharedShreddingFieldMeta result;

    // field_dict (compressed)
    PAIMON_ASSIGN_OR_RAISE(
        int32_t original_len,
        GetRequiredInt32(metadata, MapSharedShreddingDefine::kFieldDictOriginalSize));
    PAIMON_ASSIGN_OR_RAISE(std::string compressed_dict,
                           GetRequiredValue(metadata, MapSharedShreddingDefine::kFieldDict));
    int32_t compression_index = metadata->FindKey(MapSharedShreddingDefine::kFieldDictCompression);
    std::string field_dict_compression = compression_index < 0
                                             ? MapSharedShreddingDefine::kDefaultDictCompression
                                             : metadata->value(compression_index);
    field_dict_compression = StringUtils::ToLowerCase(field_dict_compression);
    PAIMON_ASSIGN_OR_RAISE(std::string field_dict_json,
                           DecompressString(compressed_dict, original_len, field_dict_compression));
    PAIMON_ASSIGN_OR_RAISE(result.name_to_id, DeserializeFieldDict(field_dict_json));

    // field_columns
    PAIMON_ASSIGN_OR_RAISE(std::string field_columns_json,
                           GetRequiredValue(metadata, MapSharedShreddingDefine::kFieldColumns));
    PAIMON_ASSIGN_OR_RAISE(result.field_to_columns, DeserializeFieldColumns(field_columns_json));

    // overflow_set
    PAIMON_ASSIGN_OR_RAISE(std::string overflow_json,
                           GetRequiredValue(metadata, MapSharedShreddingDefine::kOverflowSet));
    PAIMON_ASSIGN_OR_RAISE(result.overflow_field_set, DeserializeOverflowSet(overflow_json));

    // num_columns & max_row_width
    PAIMON_ASSIGN_OR_RAISE(result.num_columns,
                           GetRequiredInt32(metadata, MapSharedShreddingDefine::kNumColumns));
    PAIMON_ASSIGN_OR_RAISE(result.max_row_width,
                           GetRequiredInt32(metadata, MapSharedShreddingDefine::kMaxRowWidth));

    return result;
}

bool MapSharedShreddingUtils::HasShreddingMetadata(
    const std::shared_ptr<arrow::KeyValueMetadata>& metadata) {
    if (!metadata) {
        return false;
    }
    auto index = metadata->FindKey(MapShreddingDefine::kStorageLayout);
    if (index < 0) {
        return false;
    }
    return metadata->value(index) == MapShreddingDefine::kStorageLayoutSharedShredding;
}

Result<bool> MapSharedShreddingUtils::IsOverflowField(const MapSharedShreddingFieldMeta& meta,
                                                      const std::string& name) {
    auto name_iter = meta.name_to_id.find(name);
    if (name_iter == meta.name_to_id.end()) {
        return Status::Invalid(
            fmt::format("cannot find field {} in map shared shredding meta", name));
    }
    return meta.overflow_field_set.count(name_iter->second) > 0;
}

std::function<Result<std::shared_ptr<arrow::Schema>>()>
MapSharedShreddingUtils::BuildMetadataFinalizer(
    const std::shared_ptr<MapSharedShreddingBatchConverter>& converter,
    const std::string& compression, const std::shared_ptr<arrow::Schema>& physical_schema) {
    return [converter, compression, physical_schema]() -> Result<std::shared_ptr<arrow::Schema>> {
        const std::vector<std::string>& shredding_field_names =
            converter->GetShreddingColumnNames();
        arrow::FieldVector updated_fields = physical_schema->fields();
        for (const std::string& field_name : shredding_field_names) {
            int32_t col_index = physical_schema->GetFieldIndex(field_name);
            if (col_index < 0) {
                return Status::Invalid(fmt::format(
                    "Shared-shredding field '{}' not found in physical schema.", field_name));
            }
            const auto& field = physical_schema->field(col_index);
            auto metadata = field->metadata() ? field->metadata()->Copy()
                                              : std::make_shared<arrow::KeyValueMetadata>();
            PAIMON_ASSIGN_OR_RAISE(MapSharedShreddingFieldMeta file_meta,
                                   converter->BuildFieldMeta(field_name));
            PAIMON_RETURN_NOT_OK(
                MapSharedShreddingUtils::SerializeMetadata(file_meta, compression, metadata.get()));
            updated_fields[col_index] = field->WithMetadata(metadata);
        }
        return arrow::schema(std::move(updated_fields));
    };
}

}  // namespace paimon
