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

package org.apache.paimon.arrow;

import org.apache.paimon.format.FormatMetadataUtils;
import org.apache.paimon.types.DataField;
import org.apache.paimon.types.DataTypes;
import org.apache.paimon.types.RowType;

import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.FieldType;
import org.apache.arrow.vector.types.pojo.Schema;
import org.junit.jupiter.api.Test;

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

import static org.assertj.core.api.Assertions.assertThat;

/** Compatibility tests between Arrow Java schema serialization and format metadata utilities. */
public class ArrowSchemaMetadataCompatibilityTest {

    @Test
    public void testFormatMetadataCanBeReadByArrowJava() {
        RowType rowType = rowType();
        Map<String, String> tagsMetadata = tagsMetadata();
        Map<String, Map<String, String>> fieldMetadata = new LinkedHashMap<>();
        fieldMetadata.put("tags", tagsMetadata);

        byte[] schemaBytes =
                FormatMetadataUtils.buildArrowSchemaMetadata(
                        rowType, fieldMetadata, FormatMetadataUtils.PARQUET_FIELD_ID_KEY);
        Schema schema = Schema.deserializeMessage(ByteBuffer.wrap(schemaBytes));

        assertThat(schema.getFields()).extracting(Field::getName).containsExactly("id", "tags");
        assertThat(schema.findField("id").getMetadata())
                .containsEntry(ArrowUtils.PARQUET_FIELD_ID, "0");
        assertThat(schema.findField("tags").getMetadata())
                .containsEntry(ArrowUtils.PARQUET_FIELD_ID, "1")
                .containsAllEntriesOf(tagsMetadata);
    }

    @Test
    public void testRichFormatMetadataSchemaMatchesArrowUtils() {
        // Keep this case broad because Arrow schema metadata is used as an IPC-compatible
        // format contract. It verifies that format-generated bytes deserialize to the same
        // field tree as ArrowUtils, and that Arrow Java official serialization can be parsed
        // back by format metadata utilities with the same top-level field metadata.
        RowType rowType = richRowType();
        Map<String, String> nestedMetadata = new LinkedHashMap<>();
        nestedMetadata.put("paimon.test.nested", "enabled");
        Map<String, Map<String, String>> fieldMetadata = new LinkedHashMap<>();
        fieldMetadata.put("nested", nestedMetadata);

        byte[] schemaBytes =
                FormatMetadataUtils.buildArrowSchemaMetadata(
                        rowType, fieldMetadata, FormatMetadataUtils.PARQUET_FIELD_ID_KEY);
        Schema schema = Schema.deserializeMessage(ByteBuffer.wrap(schemaBytes));
        List<Field> expectedFields = arrowFields(rowType, fieldMetadata);

        assertThat(schema.getFields()).hasSize(expectedFields.size());
        for (int i = 0; i < expectedFields.size(); i++) {
            assertFieldEquals(schema.getFields().get(i), expectedFields.get(i));
        }

        byte[] arrowJavaSchemaBytes = new Schema(expectedFields).serializeAsMessage();
        assertThat(FormatMetadataUtils.readFieldMetadata(arrowJavaSchemaBytes))
                .isEqualTo(FormatMetadataUtils.readFieldMetadata(schemaBytes));
    }

    @Test
    public void testArrowJavaSchemaCanBeReadByFormatMetadata() {
        RowType rowType = rowType();
        Map<String, String> tagsMetadata = tagsMetadata();
        Map<String, Map<String, String>> fieldMetadata = new LinkedHashMap<>();
        fieldMetadata.put("tags", tagsMetadata);
        List<Field> fields = arrowFields(rowType, fieldMetadata);
        byte[] schemaBytes = new Schema(fields).serializeAsMessage();

        Map<String, Map<String, String>> metadata =
                FormatMetadataUtils.readFieldMetadata(schemaBytes);

        assertThat(metadata).containsOnlyKeys("id", "tags");
        assertThat(metadata.get("id")).containsEntry(ArrowUtils.PARQUET_FIELD_ID, "0");
        assertThat(metadata.get("tags"))
                .containsEntry(ArrowUtils.PARQUET_FIELD_ID, "1")
                .containsAllEntriesOf(tagsMetadata);
    }

    private static RowType rowType() {
        return DataTypes.ROW(
                DataTypes.FIELD(0, "id", DataTypes.INT()),
                DataTypes.FIELD(1, "tags", DataTypes.MAP(DataTypes.STRING(), DataTypes.INT())));
    }

    private static RowType richRowType() {
        return DataTypes.ROW(
                DataTypes.FIELD(0, "boolean_field", DataTypes.BOOLEAN().notNull()),
                DataTypes.FIELD(1, "tinyint_field", DataTypes.TINYINT()),
                DataTypes.FIELD(2, "smallint_field", DataTypes.SMALLINT()),
                DataTypes.FIELD(3, "int_field", DataTypes.INT()),
                DataTypes.FIELD(4, "bigint_field", DataTypes.BIGINT()),
                DataTypes.FIELD(5, "float_field", DataTypes.FLOAT()),
                DataTypes.FIELD(6, "double_field", DataTypes.DOUBLE()),
                DataTypes.FIELD(7, "decimal_field", DataTypes.DECIMAL(20, 3)),
                DataTypes.FIELD(8, "date_field", DataTypes.DATE()),
                DataTypes.FIELD(9, "time_field", DataTypes.TIME(6)),
                DataTypes.FIELD(10, "timestamp_second_field", DataTypes.TIMESTAMP(0)),
                DataTypes.FIELD(11, "timestamp_milli_field", DataTypes.TIMESTAMP(3)),
                DataTypes.FIELD(12, "timestamp_micro_field", DataTypes.TIMESTAMP(6)),
                DataTypes.FIELD(13, "timestamp_nano_field", DataTypes.TIMESTAMP(9)),
                DataTypes.FIELD(
                        14,
                        "timestamp_ltz_second_field",
                        DataTypes.TIMESTAMP_WITH_LOCAL_TIME_ZONE(0)),
                DataTypes.FIELD(
                        15,
                        "timestamp_ltz_milli_field",
                        DataTypes.TIMESTAMP_WITH_LOCAL_TIME_ZONE(3)),
                DataTypes.FIELD(
                        16,
                        "timestamp_ltz_micro_field",
                        DataTypes.TIMESTAMP_WITH_LOCAL_TIME_ZONE(6)),
                DataTypes.FIELD(
                        17,
                        "timestamp_ltz_nano_field",
                        DataTypes.TIMESTAMP_WITH_LOCAL_TIME_ZONE(9)),
                DataTypes.FIELD(18, "char_field", DataTypes.CHAR(10)),
                DataTypes.FIELD(19, "varchar_field", DataTypes.VARCHAR(20)),
                DataTypes.FIELD(20, "binary_field", DataTypes.BINARY(16)),
                DataTypes.FIELD(21, "varbinary_field", DataTypes.VARBINARY(32)),
                DataTypes.FIELD(22, "array_field", DataTypes.ARRAY(DataTypes.INT())),
                DataTypes.FIELD(
                        23,
                        "map_field",
                        DataTypes.MAP(
                                DataTypes.STRING(), DataTypes.ARRAY(DataTypes.DECIMAL(10, 2)))),
                DataTypes.FIELD(
                        24,
                        "nested",
                        DataTypes.ROW(
                                DataTypes.FIELD(25, "name", DataTypes.STRING()),
                                DataTypes.FIELD(
                                        26,
                                        "scores",
                                        DataTypes.ARRAY(
                                                DataTypes.ROW(
                                                        DataTypes.FIELD(
                                                                27, "score", DataTypes.INT())))))),
                DataTypes.FIELD(28, "vector_field", DataTypes.VECTOR(3, DataTypes.FLOAT())),
                DataTypes.FIELD(29, "variant_field", DataTypes.VARIANT()));
    }

    private static Map<String, String> tagsMetadata() {
        Map<String, String> metadata = new LinkedHashMap<>();
        metadata.put("paimon.test.tags", "enabled");
        metadata.put("paimon.test.version", "1");
        return metadata;
    }

    private static List<Field> arrowFields(
            RowType rowType, Map<String, Map<String, String>> fieldMetadata) {
        List<Field> fields = new ArrayList<>();
        for (DataField field : rowType.getFields()) {
            Field arrowField = ArrowUtils.toArrowField(field.name(), field.id(), field.type(), 0);
            Map<String, String> metadata = fieldMetadata.get(field.name());
            if (metadata != null) {
                arrowField = withMetadata(arrowField, metadata);
            }
            fields.add(arrowField);
        }
        return fields;
    }

    private static Field withMetadata(Field field, Map<String, String> metadata) {
        Map<String, String> merged = new LinkedHashMap<>();
        merged.putAll(metadata);
        merged.putAll(field.getMetadata());
        FieldType fieldType = field.getFieldType();
        return new Field(
                field.getName(),
                new FieldType(
                        fieldType.isNullable(),
                        fieldType.getType(),
                        fieldType.getDictionary(),
                        merged),
                field.getChildren());
    }

    private static void assertFieldEquals(Field actual, Field expected) {
        assertThat(actual.getName()).isEqualTo(expected.getName());
        assertThat(actual.getFieldType().isNullable())
                .isEqualTo(expected.getFieldType().isNullable());
        assertThat(actual.getFieldType().getType()).isEqualTo(expected.getFieldType().getType());
        assertThat(actual.getFieldType().getDictionary())
                .isEqualTo(expected.getFieldType().getDictionary());
        assertThat(actual.getMetadata()).isEqualTo(expected.getMetadata());
        assertThat(actual.getChildren()).hasSize(expected.getChildren().size());
        for (int i = 0; i < expected.getChildren().size(); i++) {
            assertFieldEquals(actual.getChildren().get(i), expected.getChildren().get(i));
        }
    }
}
