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

/* This file is based on source code from the Spark Project (http://spark.apache.org/), licensed
 * by the Apache Software Foundation (ASF) under the Apache License, Version 2.0. See the NOTICE
 * file distributed with this work for additional information regarding copyright ownership. */

#include "paimon/common/data/variant/variant_shredding_utils.h"

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

#include "arrow/api.h"
#include "arrow/util/checked_cast.h"
#include "fmt/format.h"
#include "paimon/common/data/variant/variant_defs.h"
#include "paimon/common/data/variant/variant_type_utils.h"

namespace paimon {

namespace {

Status InvalidVariantShreddingSchema(const std::shared_ptr<arrow::DataType>& type) {
    return Status::Invalid(
        fmt::format("Invalid variant shredding schema: {}", type ? type->ToString() : "null"));
}

// Mirrors the Java `PaimonShreddingUtils.variantShreddingSchema(dataType, isTopLevel,
// isObjectField)`.
Result<std::shared_ptr<arrow::DataType>> VariantShreddingSchemaImpl(
    const std::shared_ptr<arrow::DataType>& data_type, bool is_top_level, bool is_object_field) {
    arrow::FieldVector fields;
    if (is_top_level) {
        fields.push_back(arrow::field(VariantDefs::kMetadataFieldName, arrow::binary(),
                                      /*nullable=*/false));
    }
    switch (data_type->id()) {
        case arrow::Type::LIST: {
            const auto& list_type = std::static_pointer_cast<arrow::ListType>(data_type);
            PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::DataType> element_type,
                                   VariantShreddingSchemaImpl(list_type->value_type(),
                                                              /*is_top_level=*/false,
                                                              /*is_object_field=*/false));
            fields.push_back(
                arrow::field(VariantDefs::kValueFieldName, arrow::binary(), /*nullable=*/true));
            fields.push_back(arrow::field(VariantDefs::kTypedValueFieldName,
                                          arrow::list(element_type), /*nullable=*/true));
            break;
        }
        case arrow::Type::STRUCT: {
            // The field name level is always non-nullable: Variant null values are represented in
            // the "value" column as "00", and missing values are represented by setting both
            // "value" and "typed_value" to null.
            const auto& struct_type = std::static_pointer_cast<arrow::StructType>(data_type);
            arrow::FieldVector shredded_fields;
            for (const auto& field : struct_type->fields()) {
                PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::DataType> field_type,
                                       VariantShreddingSchemaImpl(field->type(),
                                                                  /*is_top_level=*/false,
                                                                  /*is_object_field=*/true));
                shredded_fields.push_back(
                    arrow::field(field->name(), field_type, /*nullable=*/false));
            }
            fields.push_back(
                arrow::field(VariantDefs::kValueFieldName, arrow::binary(), /*nullable=*/true));
            fields.push_back(arrow::field(VariantDefs::kTypedValueFieldName,
                                          arrow::struct_(shredded_fields), /*nullable=*/true));
            break;
        }
        case arrow::Type::NA: {
            // `arrow::null()` denotes an untyped VARIANT leaf in shredding types. It doesn't
            // need a typed column. If there is no typed column, value is required for array
            // elements or top-level fields, but optional for objects (where a null represents a
            // missing field).
            fields.push_back(arrow::field(VariantDefs::kValueFieldName, arrow::binary(),
                                          /*nullable=*/is_object_field));
            break;
        }
        case arrow::Type::STRING:
        case arrow::Type::BOOL:
        case arrow::Type::BINARY:
        case arrow::Type::DECIMAL128:
        case arrow::Type::INT8:
        case arrow::Type::INT16:
        case arrow::Type::INT32:
        case arrow::Type::INT64:
        case arrow::Type::FLOAT:
        case arrow::Type::DOUBLE: {
            fields.push_back(
                arrow::field(VariantDefs::kValueFieldName, arrow::binary(), /*nullable=*/true));
            fields.push_back(arrow::field(VariantDefs::kTypedValueFieldName, data_type,
                                          /*nullable=*/true));
            break;
        }
        default:
            return InvalidVariantShreddingSchema(data_type);
    }
    return arrow::struct_(fields);
}

Result<std::shared_ptr<VariantSchema>> BuildVariantSchemaImpl(
    const std::shared_ptr<arrow::DataType>& type, bool top_level) {
    if (type->id() != arrow::Type::STRUCT) {
        return InvalidVariantShreddingSchema(type);
    }
    const auto& struct_type = std::static_pointer_cast<arrow::StructType>(type);
    // The struct must not be empty or contain duplicate field names. The latter is enforced in
    // the loop below.
    if (struct_type->num_fields() == 0) {
        return InvalidVariantShreddingSchema(type);
    }

    auto schema = std::make_shared<VariantSchema>();
    schema->num_fields = struct_type->num_fields();

    for (int32_t i = 0; i < struct_type->num_fields(); ++i) {
        const auto& field = struct_type->field(i);
        const auto& field_type = field->type();
        if (field->name() == VariantDefs::kTypedValueFieldName) {
            if (schema->typed_idx != -1) {
                return InvalidVariantShreddingSchema(type);
            }
            schema->typed_idx = i;
            switch (field_type->id()) {
                case arrow::Type::STRUCT: {
                    const auto& object_type =
                        std::static_pointer_cast<arrow::StructType>(field_type);
                    schema->has_object_schema = true;
                    schema->object_schema.reserve(object_type->num_fields());
                    for (int32_t index = 0; index < object_type->num_fields(); ++index) {
                        const auto& object_field = object_type->field(index);
                        PAIMON_ASSIGN_OR_RAISE(
                            std::shared_ptr<VariantSchema> field_schema,
                            BuildVariantSchemaImpl(object_field->type(), /*top_level=*/false));
                        schema->object_schema.push_back(
                            VariantSchema::ObjectField{object_field->name(), field_schema});
                        auto [it, inserted] =
                            schema->object_schema_map.emplace(object_field->name(), index);
                        if (!inserted) {
                            return InvalidVariantShreddingSchema(type);
                        }
                    }
                    break;
                }
                case arrow::Type::LIST: {
                    const auto& list_type = std::static_pointer_cast<arrow::ListType>(field_type);
                    PAIMON_ASSIGN_OR_RAISE(
                        schema->array_schema,
                        BuildVariantSchemaImpl(list_type->value_type(), /*top_level=*/false));
                    break;
                }
                case arrow::Type::BOOL:
                    schema->scalar_schema =
                        VariantSchema::ScalarType{VariantSchema::ScalarKind::kBoolean};
                    break;
                case arrow::Type::INT8:
                    schema->scalar_schema =
                        VariantSchema::ScalarType{VariantSchema::ScalarKind::kByte};
                    break;
                case arrow::Type::INT16:
                    schema->scalar_schema =
                        VariantSchema::ScalarType{VariantSchema::ScalarKind::kShort};
                    break;
                case arrow::Type::INT32:
                    schema->scalar_schema =
                        VariantSchema::ScalarType{VariantSchema::ScalarKind::kInt};
                    break;
                case arrow::Type::INT64:
                    schema->scalar_schema =
                        VariantSchema::ScalarType{VariantSchema::ScalarKind::kLong};
                    break;
                case arrow::Type::FLOAT:
                    schema->scalar_schema =
                        VariantSchema::ScalarType{VariantSchema::ScalarKind::kFloat};
                    break;
                case arrow::Type::DOUBLE:
                    schema->scalar_schema =
                        VariantSchema::ScalarType{VariantSchema::ScalarKind::kDouble};
                    break;
                case arrow::Type::STRING:
                    schema->scalar_schema =
                        VariantSchema::ScalarType{VariantSchema::ScalarKind::kString};
                    break;
                case arrow::Type::BINARY:
                    schema->scalar_schema =
                        VariantSchema::ScalarType{VariantSchema::ScalarKind::kBinary};
                    break;
                case arrow::Type::DATE32:
                    schema->scalar_schema =
                        VariantSchema::ScalarType{VariantSchema::ScalarKind::kDate};
                    break;
                case arrow::Type::DECIMAL128: {
                    const auto& decimal_type =
                        std::static_pointer_cast<arrow::Decimal128Type>(field_type);
                    schema->scalar_schema =
                        VariantSchema::ScalarType{VariantSchema::ScalarKind::kDecimal,
                                                  decimal_type->precision(), decimal_type->scale()};
                    break;
                }
                case arrow::Type::TIMESTAMP: {
                    const auto& timestamp_type =
                        std::static_pointer_cast<arrow::TimestampType>(field_type);
                    // The variant binary stores timestamps as microseconds since the epoch; a
                    // typed_value column of any other precision would misinterpret the values.
                    if (timestamp_type->unit() != arrow::TimeUnit::MICRO) {
                        return InvalidVariantShreddingSchema(type);
                    }
                    schema->scalar_schema =
                        VariantSchema::ScalarType{timestamp_type->timezone().empty()
                                                      ? VariantSchema::ScalarKind::kTimestampNtz
                                                      : VariantSchema::ScalarKind::kTimestampLtz};
                    break;
                }
                default:
                    return InvalidVariantShreddingSchema(type);
            }
        } else if (field->name() == VariantDefs::kValueFieldName) {
            if (schema->variant_idx != -1 || field_type->id() != arrow::Type::BINARY) {
                return InvalidVariantShreddingSchema(type);
            }
            schema->variant_idx = i;
        } else if (field->name() == VariantDefs::kMetadataFieldName) {
            if (schema->top_level_metadata_idx != -1 || field_type->id() != arrow::Type::BINARY) {
                return InvalidVariantShreddingSchema(type);
            }
            schema->top_level_metadata_idx = i;
        } else {
            return InvalidVariantShreddingSchema(type);
        }
    }

    if (top_level != (schema->top_level_metadata_idx >= 0)) {
        return InvalidVariantShreddingSchema(type);
    }
    return schema;
}

}  // namespace

Result<std::shared_ptr<arrow::DataType>> VariantShreddingUtils::VariantShreddingSchema(
    const std::shared_ptr<arrow::DataType>& shredding_type) {
    return VariantShreddingSchemaImpl(shredding_type, /*is_top_level=*/true,
                                      /*is_object_field=*/false);
}

Result<std::shared_ptr<VariantSchema>> VariantShreddingUtils::BuildVariantSchema(
    const std::shared_ptr<arrow::DataType>& struct_type) {
    return BuildVariantSchemaImpl(struct_type, /*top_level=*/true);
}

Result<std::shared_ptr<arrow::DataType>> VariantShreddingUtils::ScalarSchemaToArrowType(
    const VariantSchema::ScalarType& scalar) {
    switch (scalar.kind) {
        case VariantSchema::ScalarKind::kBoolean:
            return arrow::boolean();
        case VariantSchema::ScalarKind::kByte:
            return arrow::int8();
        case VariantSchema::ScalarKind::kShort:
            return arrow::int16();
        case VariantSchema::ScalarKind::kInt:
            return arrow::int32();
        case VariantSchema::ScalarKind::kLong:
            return arrow::int64();
        case VariantSchema::ScalarKind::kFloat:
            return arrow::float32();
        case VariantSchema::ScalarKind::kDouble:
            return arrow::float64();
        case VariantSchema::ScalarKind::kString:
            return arrow::utf8();
        case VariantSchema::ScalarKind::kBinary:
            return arrow::binary();
        case VariantSchema::ScalarKind::kDecimal:
            return arrow::decimal128(scalar.precision, scalar.scale);
        case VariantSchema::ScalarKind::kDate:
            return arrow::date32();
        case VariantSchema::ScalarKind::kTimestampLtz:
            return arrow::timestamp(arrow::TimeUnit::MICRO, "UTC");
        case VariantSchema::ScalarKind::kTimestampNtz:
            return arrow::timestamp(arrow::TimeUnit::MICRO);
        default:
            return Status::NotImplemented(fmt::format("Unsupported variant scalar kind: {}",
                                                      static_cast<int32_t>(scalar.kind)));
    }
}

bool VariantShreddingUtils::IsShreddedFileType(
    const std::shared_ptr<arrow::DataType>& file_variant_type) {
    if (!file_variant_type || file_variant_type->id() != arrow::Type::STRUCT) {
        return false;
    }
    const auto& struct_type = std::static_pointer_cast<arrow::StructType>(file_variant_type);
    return struct_type->GetFieldByName(VariantDefs::kTypedValueFieldName) != nullptr;
}

bool VariantShreddingUtils::IsUntypedPhysicalVariantType(
    const std::shared_ptr<arrow::DataType>& file_variant_type) {
    if (!file_variant_type || file_variant_type->id() != arrow::Type::STRUCT) {
        return false;
    }
    const auto& struct_type = std::static_pointer_cast<arrow::StructType>(file_variant_type);
    if (struct_type->num_fields() != 2) {
        return false;
    }
    const auto& metadata = struct_type->field(0);
    const auto& value = struct_type->field(1);
    return metadata->name() == VariantDefs::kMetadataFieldName &&
           metadata->type()->id() == arrow::Type::BINARY &&
           value->name() == VariantDefs::kValueFieldName &&
           value->type()->id() == arrow::Type::BINARY;
}

}  // namespace paimon
