/*
 * 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.iceberg.types;

import static org.apache.iceberg.types.Types.NestedField.optional;
import static org.apache.iceberg.types.Types.NestedField.required;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;

import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Stream;
import org.apache.iceberg.Schema;
import org.apache.iceberg.expressions.Literal;
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
import org.apache.iceberg.relocated.com.google.common.collect.Lists;
import org.apache.iceberg.relocated.com.google.common.collect.Sets;
import org.apache.iceberg.types.Types.IntegerType;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;

public class TestTypeUtil {
  @Test
  public void testReassignIdsDuplicateColumns() {
    Schema schema =
        new Schema(
            required(0, "a", Types.IntegerType.get()), required(1, "A", Types.IntegerType.get()));
    Schema sourceSchema =
        new Schema(
            required(1, "a", Types.IntegerType.get()), required(2, "A", Types.IntegerType.get()));
    final Schema actualSchema = TypeUtil.reassignIds(schema, sourceSchema);
    assertThat(actualSchema.asStruct()).isEqualTo(sourceSchema.asStruct());
  }

  @Test
  public void testReassignIdsWithIdentifier() {
    Schema schema =
        new Schema(
            Lists.newArrayList(
                required(0, "a", Types.IntegerType.get()),
                required(1, "A", Types.IntegerType.get())),
            Sets.newHashSet(0));
    Schema sourceSchema =
        new Schema(
            Lists.newArrayList(
                required(1, "a", Types.IntegerType.get()),
                required(2, "A", Types.IntegerType.get())),
            Sets.newHashSet(1));
    final Schema actualSchema = TypeUtil.reassignIds(schema, sourceSchema);
    assertThat(actualSchema.asStruct()).isEqualTo(sourceSchema.asStruct());
    assertThat(actualSchema.identifierFieldIds())
        .as("identifier field ID should change based on source schema")
        .isEqualTo(sourceSchema.identifierFieldIds());
  }

  @Test
  public void testAssignIncreasingFreshIdWithIdentifier() {
    Schema schema =
        new Schema(
            Lists.newArrayList(
                required(10, "a", Types.IntegerType.get()),
                required(11, "A", Types.IntegerType.get())),
            Sets.newHashSet(10));
    Schema expectedSchema =
        new Schema(
            Lists.newArrayList(
                required(1, "a", Types.IntegerType.get()),
                required(2, "A", Types.IntegerType.get())),
            Sets.newHashSet(1));
    final Schema actualSchema = TypeUtil.assignIncreasingFreshIds(schema);
    assertThat(actualSchema.asStruct()).isEqualTo(expectedSchema.asStruct());
    assertThat(actualSchema.identifierFieldIds())
        .as("identifier field ID should change based on source schema")
        .isEqualTo(expectedSchema.identifierFieldIds());
  }

  @Test
  public void testAssignIncreasingFreshIdNewIdentifier() {
    Schema schema =
        new Schema(
            Lists.newArrayList(
                required(10, "a", Types.IntegerType.get()),
                required(11, "A", Types.IntegerType.get())),
            Sets.newHashSet(10));
    Schema sourceSchema =
        new Schema(
            Lists.newArrayList(
                required(1, "a", Types.IntegerType.get()),
                required(2, "A", Types.IntegerType.get())));
    final Schema actualSchema = TypeUtil.reassignIds(schema, sourceSchema);
    assertThat(actualSchema.asStruct()).isEqualTo(sourceSchema.asStruct());
    assertThat(actualSchema.identifierFieldIds())
        .as("source schema missing identifier should not impact refreshing new identifier")
        .isEqualTo(Sets.newHashSet(sourceSchema.findField("a").fieldId()));
  }

  @Test
  public void testProject() {
    Schema schema =
        new Schema(
            Lists.newArrayList(
                required(10, "a", Types.IntegerType.get()),
                required(11, "A", Types.IntegerType.get()),
                required(
                    12,
                    "someStruct",
                    Types.StructType.of(
                        required(13, "b", Types.IntegerType.get()),
                        required(14, "B", Types.IntegerType.get()),
                        required(
                            15,
                            "anotherStruct",
                            Types.StructType.of(
                                required(16, "c", Types.IntegerType.get()),
                                required(17, "C", Types.IntegerType.get())))))));

    Schema expectedTop = new Schema(Lists.newArrayList(required(11, "A", Types.IntegerType.get())));

    Schema actualTop = TypeUtil.project(schema, Sets.newHashSet(11));
    assertThat(actualTop.asStruct()).isEqualTo(expectedTop.asStruct());

    Schema expectedDepthOne =
        new Schema(
            Lists.newArrayList(
                required(10, "a", Types.IntegerType.get()),
                required(
                    12,
                    "someStruct",
                    Types.StructType.of(required(13, "b", Types.IntegerType.get())))));

    Schema actualDepthOne = TypeUtil.project(schema, Sets.newHashSet(10, 12, 13));
    assertThat(actualDepthOne.asStruct()).isEqualTo(expectedDepthOne.asStruct());

    Schema expectedDepthTwo =
        new Schema(
            Lists.newArrayList(
                required(11, "A", Types.IntegerType.get()),
                required(
                    12,
                    "someStruct",
                    Types.StructType.of(
                        required(
                            15,
                            "anotherStruct",
                            Types.StructType.of(required(17, "C", Types.IntegerType.get())))))));

    Schema actualDepthTwo = TypeUtil.project(schema, Sets.newHashSet(11, 12, 15, 17));
    Schema actualDepthTwoChildren = TypeUtil.project(schema, Sets.newHashSet(11, 17));
    assertThat(actualDepthTwo.asStruct()).isEqualTo(expectedDepthTwo.asStruct());
    assertThat(actualDepthTwoChildren.asStruct()).isEqualTo(expectedDepthTwo.asStruct());
  }

  @Test
  public void testProjectNaturallyEmpty() {
    Schema schema =
        new Schema(
            Lists.newArrayList(
                required(
                    12,
                    "someStruct",
                    Types.StructType.of(
                        required(
                            15,
                            "anotherStruct",
                            Types.StructType.of(required(20, "empty", Types.StructType.of())))))));

    Schema expectedDepthOne =
        new Schema(Lists.newArrayList(required(12, "someStruct", Types.StructType.of())));

    Schema actualDepthOne = TypeUtil.project(schema, Sets.newHashSet(12));
    assertThat(actualDepthOne.asStruct()).isEqualTo(expectedDepthOne.asStruct());

    Schema expectedDepthTwo =
        new Schema(
            Lists.newArrayList(
                required(
                    12,
                    "someStruct",
                    Types.StructType.of(required(15, "anotherStruct", Types.StructType.of())))));

    Schema actualDepthTwo = TypeUtil.project(schema, Sets.newHashSet(12, 15));
    assertThat(actualDepthTwo.asStruct()).isEqualTo(expectedDepthTwo.asStruct());

    Schema expectedDepthThree =
        new Schema(
            Lists.newArrayList(
                required(
                    12,
                    "someStruct",
                    Types.StructType.of(
                        required(
                            15,
                            "anotherStruct",
                            Types.StructType.of(required(20, "empty", Types.StructType.of())))))));

    Schema actualDepthThree = TypeUtil.project(schema, Sets.newHashSet(12, 15, 20));
    Schema actualDepthThreeChildren = TypeUtil.project(schema, Sets.newHashSet(20));
    assertThat(actualDepthThree.asStruct()).isEqualTo(expectedDepthThree.asStruct());
    assertThat(actualDepthThreeChildren.asStruct()).isEqualTo(expectedDepthThree.asStruct());
  }

  @Test
  public void testProjectEmpty() {
    Schema schema =
        new Schema(
            Lists.newArrayList(
                required(10, "a", Types.IntegerType.get()),
                required(11, "A", Types.IntegerType.get()),
                required(
                    12,
                    "someStruct",
                    Types.StructType.of(
                        required(13, "b", Types.IntegerType.get()),
                        required(14, "B", Types.IntegerType.get()),
                        required(
                            15,
                            "anotherStruct",
                            Types.StructType.of(
                                required(16, "c", Types.IntegerType.get()),
                                required(17, "C", Types.IntegerType.get())))))));

    Schema expectedDepthOne =
        new Schema(Lists.newArrayList(required(12, "someStruct", Types.StructType.of())));

    Schema actualDepthOne = TypeUtil.project(schema, Sets.newHashSet(12));
    assertThat(actualDepthOne.asStruct()).isEqualTo(expectedDepthOne.asStruct());

    Schema expectedDepthTwo =
        new Schema(
            Lists.newArrayList(
                required(
                    12,
                    "someStruct",
                    Types.StructType.of(required(15, "anotherStruct", Types.StructType.of())))));

    Schema actualDepthTwo = TypeUtil.project(schema, Sets.newHashSet(12, 15));
    assertThat(actualDepthTwo.asStruct()).isEqualTo(expectedDepthTwo.asStruct());
  }

  @Test
  public void testSelect() {
    Schema schema =
        new Schema(
            Lists.newArrayList(
                required(10, "a", Types.IntegerType.get()),
                required(11, "A", Types.IntegerType.get()),
                required(
                    12,
                    "someStruct",
                    Types.StructType.of(
                        required(13, "b", Types.IntegerType.get()),
                        required(14, "B", Types.IntegerType.get()),
                        required(
                            15,
                            "anotherStruct",
                            Types.StructType.of(
                                required(16, "c", Types.IntegerType.get()),
                                required(17, "C", Types.IntegerType.get())))))));

    Schema expectedTop = new Schema(Lists.newArrayList(required(11, "A", Types.IntegerType.get())));

    Schema actualTop = TypeUtil.select(schema, Sets.newHashSet(11));
    assertThat(actualTop.asStruct()).isEqualTo(expectedTop.asStruct());

    Schema expectedDepthOne =
        new Schema(
            Lists.newArrayList(
                required(10, "a", Types.IntegerType.get()),
                required(
                    12,
                    "someStruct",
                    Types.StructType.of(
                        required(13, "b", Types.IntegerType.get()),
                        required(14, "B", Types.IntegerType.get()),
                        required(
                            15,
                            "anotherStruct",
                            Types.StructType.of(
                                required(16, "c", Types.IntegerType.get()),
                                required(17, "C", Types.IntegerType.get())))))));

    Schema actualDepthOne = TypeUtil.select(schema, Sets.newHashSet(10, 12));
    assertThat(actualDepthOne.asStruct()).isEqualTo(expectedDepthOne.asStruct());

    Schema expectedDepthTwo =
        new Schema(
            Lists.newArrayList(
                required(11, "A", Types.IntegerType.get()),
                required(
                    12,
                    "someStruct",
                    Types.StructType.of(
                        required(
                            15,
                            "anotherStruct",
                            Types.StructType.of(required(17, "C", Types.IntegerType.get())))))));

    Schema actualDepthTwo = TypeUtil.select(schema, Sets.newHashSet(11, 17));
    assertThat(actualDepthTwo.asStruct()).isEqualTo(expectedDepthTwo.asStruct());
  }

  @Test
  public void testSelectInIdOrder() {
    Schema schema =
        new Schema(
            required(1, "id", Types.IntegerType.get()),
            required(3, "b", Types.IntegerType.get()),
            required(2, "a", Types.IntegerType.get()));

    Schema result = TypeUtil.selectInIdOrder(schema, Sets.newHashSet(2, 3));

    assertThat(result.columns()).hasSize(2);
    assertThat(result.columns().get(0).fieldId()).isEqualTo(2);
    assertThat(result.columns().get(1).fieldId()).isEqualTo(3);

    // verify that different input orderings produce the same result
    Schema schemaReversed =
        new Schema(
            required(2, "a", Types.IntegerType.get()),
            required(3, "b", Types.IntegerType.get()),
            required(1, "id", Types.IntegerType.get()));

    Schema resultReversed = TypeUtil.selectInIdOrder(schemaReversed, Sets.newHashSet(2, 3));
    assertThat(resultReversed.asStruct()).isEqualTo(result.asStruct());
  }

  @Test
  public void testProjectMap() {
    // We can't partially project keys because it changes key equality
    Schema schema =
        new Schema(
            Lists.newArrayList(
                required(10, "a", Types.IntegerType.get()),
                required(11, "A", Types.IntegerType.get()),
                required(
                    12,
                    "map",
                    Types.MapType.ofRequired(
                        13,
                        14,
                        Types.StructType.of(
                            optional(100, "x", Types.IntegerType.get()),
                            optional(101, "y", Types.IntegerType.get())),
                        Types.StructType.of(
                            required(200, "z", Types.IntegerType.get()),
                            optional(
                                201,
                                "innerMap",
                                Types.MapType.ofOptional(
                                    202,
                                    203,
                                    Types.IntegerType.get(),
                                    Types.StructType.of(
                                        required(300, "foo", Types.IntegerType.get()),
                                        required(301, "bar", Types.IntegerType.get())))))))));

    assertThatThrownBy(() -> TypeUtil.project(schema, Sets.newHashSet(12)))
        .isInstanceOf(IllegalArgumentException.class)
        .hasMessageContaining("Cannot explicitly project List or Map types");

    assertThatThrownBy(() -> TypeUtil.project(schema, Sets.newHashSet(201)))
        .isInstanceOf(IllegalArgumentException.class)
        .hasMessageContaining("Cannot explicitly project List or Map types");

    Schema expectedTopLevel =
        new Schema(Lists.newArrayList(required(10, "a", Types.IntegerType.get())));
    Schema actualTopLevel = TypeUtil.project(schema, Sets.newHashSet(10));
    assertThat(actualTopLevel.asStruct()).isEqualTo(expectedTopLevel.asStruct());

    Schema expectedDepthOne =
        new Schema(
            Lists.newArrayList(
                required(10, "a", Types.IntegerType.get()),
                required(
                    12,
                    "map",
                    Types.MapType.ofRequired(
                        13,
                        14,
                        Types.StructType.of(
                            optional(100, "x", Types.IntegerType.get()),
                            optional(101, "y", Types.IntegerType.get())),
                        Types.StructType.of()))));
    Schema actualDepthOne = TypeUtil.project(schema, Sets.newHashSet(10, 13, 14, 100, 101));
    Schema actualDepthOneNoKeys = TypeUtil.project(schema, Sets.newHashSet(10, 13, 14));
    assertThat(actualDepthOne.asStruct()).isEqualTo(expectedDepthOne.asStruct());
    assertThat(actualDepthOneNoKeys.asStruct()).isEqualTo(expectedDepthOne.asStruct());

    Schema expectedDepthTwo =
        new Schema(
            Lists.newArrayList(
                required(10, "a", Types.IntegerType.get()),
                required(
                    12,
                    "map",
                    Types.MapType.ofRequired(
                        13,
                        14,
                        Types.StructType.of(
                            optional(100, "x", Types.IntegerType.get()),
                            optional(101, "y", Types.IntegerType.get())),
                        Types.StructType.of(
                            required(200, "z", Types.IntegerType.get()),
                            optional(
                                201,
                                "innerMap",
                                Types.MapType.ofOptional(
                                    202, 203, Types.IntegerType.get(), Types.StructType.of())))))));
    Schema actualDepthTwo =
        TypeUtil.project(schema, Sets.newHashSet(10, 13, 14, 100, 101, 200, 202, 203));
    assertThat(actualDepthTwo.asStruct()).isEqualTo(expectedDepthTwo.asStruct());
  }

  @Test
  public void testGetProjectedIds() {
    Schema schema =
        new Schema(
            Lists.newArrayList(
                required(10, "a", Types.IntegerType.get()),
                required(11, "A", Types.IntegerType.get()),
                required(35, "emptyStruct", Types.StructType.of()),
                required(
                    12,
                    "someStruct",
                    Types.StructType.of(
                        required(13, "b", Types.IntegerType.get()),
                        required(14, "B", Types.IntegerType.get()),
                        required(
                            15,
                            "anotherStruct",
                            Types.StructType.of(
                                required(16, "c", Types.IntegerType.get()),
                                required(17, "C", Types.IntegerType.get())))))));

    Set<Integer> expectedIds = Sets.newHashSet(10, 11, 35, 12, 13, 14, 15, 16, 17);
    Set<Integer> actualIds = TypeUtil.getProjectedIds(schema);

    assertThat(actualIds).isEqualTo(expectedIds);
  }

  @Test
  public void testProjectListNested() {
    Schema schema =
        new Schema(
            Lists.newArrayList(
                required(
                    12,
                    "list",
                    Types.ListType.ofRequired(
                        13,
                        Types.ListType.ofRequired(
                            14,
                            Types.MapType.ofRequired(
                                15,
                                16,
                                IntegerType.get(),
                                Types.StructType.of(
                                    required(17, "x", Types.IntegerType.get()),
                                    required(18, "y", Types.IntegerType.get()))))))));

    assertThatThrownBy(() -> TypeUtil.project(schema, Sets.newHashSet(12)))
        .isInstanceOf(IllegalArgumentException.class)
        .hasMessageContaining("Cannot explicitly project List or Map types");

    assertThatThrownBy(() -> TypeUtil.project(schema, Sets.newHashSet(13)))
        .isInstanceOf(IllegalArgumentException.class)
        .hasMessageContaining("Cannot explicitly project List or Map types");

    assertThatThrownBy(() -> TypeUtil.project(schema, Sets.newHashSet(14)))
        .isInstanceOf(IllegalArgumentException.class)
        .hasMessageContaining("Cannot explicitly project List or Map types");

    Schema expected =
        new Schema(
            Lists.newArrayList(
                required(
                    12,
                    "list",
                    Types.ListType.ofRequired(
                        13,
                        Types.ListType.ofRequired(
                            14,
                            Types.MapType.ofRequired(
                                15, 16, IntegerType.get(), Types.StructType.of()))))));

    Schema actual = TypeUtil.project(schema, Sets.newHashSet(16));
    assertThat(actual.asStruct()).isEqualTo(expected.asStruct());
  }

  @Test
  public void testProjectMapNested() {
    Schema schema =
        new Schema(
            Lists.newArrayList(
                required(
                    12,
                    "map",
                    Types.MapType.ofRequired(
                        13,
                        14,
                        Types.IntegerType.get(),
                        Types.MapType.ofRequired(
                            15,
                            16,
                            Types.IntegerType.get(),
                            Types.ListType.ofRequired(
                                17,
                                Types.StructType.of(
                                    required(18, "x", Types.IntegerType.get()),
                                    required(19, "y", Types.IntegerType.get()))))))));

    assertThatThrownBy(() -> TypeUtil.project(schema, Sets.newHashSet(12)))
        .isInstanceOf(IllegalArgumentException.class)
        .hasMessageContaining("Cannot explicitly project List or Map types");

    assertThatThrownBy(() -> TypeUtil.project(schema, Sets.newHashSet(14)))
        .isInstanceOf(IllegalArgumentException.class)
        .hasMessageContaining("Cannot explicitly project List or Map types");

    assertThatThrownBy(() -> TypeUtil.project(schema, Sets.newHashSet(16)))
        .isInstanceOf(IllegalArgumentException.class)
        .hasMessageContaining("Cannot explicitly project List or Map types");

    Schema expected =
        new Schema(
            Lists.newArrayList(
                required(
                    12,
                    "map",
                    Types.MapType.ofRequired(
                        13,
                        14,
                        Types.IntegerType.get(),
                        Types.MapType.ofRequired(
                            15,
                            16,
                            Types.IntegerType.get(),
                            Types.ListType.ofRequired(17, Types.StructType.of()))))));

    Schema actual = TypeUtil.project(schema, Sets.newHashSet(17));
    assertThat(actual.asStruct()).isEqualTo(expected.asStruct());
  }

  @Test
  public void testReassignIdsIllegalArgumentException() {
    Schema schema =
        new Schema(
            required(1, "a", Types.IntegerType.get()), required(2, "b", Types.IntegerType.get()));
    Schema sourceSchema = new Schema(required(1, "a", Types.IntegerType.get()));
    assertThatThrownBy(() -> TypeUtil.reassignIds(schema, sourceSchema))
        .isInstanceOf(IllegalArgumentException.class)
        .hasMessage("Field b not found in source schema");
  }

  @Test
  public void testValidateSchemaViaByName() {
    Types.NestedField nestedType =
        Types.NestedField.required(
            1,
            "a",
            Types.StructType.of(
                required(2, "b", Types.StructType.of(required(3, "c", Types.BooleanType.get()))),
                required(4, "b.c", Types.BooleanType.get())));

    assertThatThrownBy(() -> TypeUtil.indexByName(Types.StructType.of(nestedType)))
        .isInstanceOf(RuntimeException.class)
        .hasMessageContaining("Invalid schema: multiple fields for name a.b.c");
  }

  @Test
  public void testSelectNot() {
    Schema schema =
        new Schema(
            Lists.newArrayList(
                required(1, "id", Types.LongType.get()),
                required(
                    2,
                    "location",
                    Types.StructType.of(
                        required(3, "lat", Types.DoubleType.get()),
                        required(4, "long", Types.DoubleType.get())))));

    Schema expectedNoPrimitive =
        new Schema(
            Lists.newArrayList(
                required(
                    2,
                    "location",
                    Types.StructType.of(
                        required(3, "lat", Types.DoubleType.get()),
                        required(4, "long", Types.DoubleType.get())))));

    Schema actualNoPrimitve = TypeUtil.selectNot(schema, Sets.newHashSet(1));
    assertThat(actualNoPrimitve.asStruct()).isEqualTo(expectedNoPrimitive.asStruct());

    // Expected legacy behavior is to completely remove structs if their elements are removed
    Schema expectedNoStructElements = new Schema(required(1, "id", Types.LongType.get()));
    Schema actualNoStructElements = TypeUtil.selectNot(schema, Sets.newHashSet(3, 4));
    assertThat(actualNoStructElements.asStruct()).isEqualTo(expectedNoStructElements.asStruct());

    // Expected legacy behavior is to ignore selectNot on struct elements.
    Schema actualNoStruct = TypeUtil.selectNot(schema, Sets.newHashSet(2));
    assertThat(actualNoStruct.asStruct()).isEqualTo(schema.asStruct());
  }

  @Test
  public void testReassignOrRefreshIds() {
    Schema schema =
        new Schema(
            Lists.newArrayList(
                required(10, "a", Types.IntegerType.get()),
                Types.NestedField.required("c")
                    .withId(11)
                    .ofType(Types.IntegerType.get())
                    .withInitialDefault(Literal.of(23))
                    .withWriteDefault(Literal.of(34))
                    .build(),
                required(12, "B", Types.IntegerType.get())),
            Sets.newHashSet(10));
    Schema sourceSchema =
        new Schema(
            Lists.newArrayList(
                required(1, "a", Types.IntegerType.get()),
                required(15, "B", Types.IntegerType.get())));

    Schema actualSchema = TypeUtil.reassignOrRefreshIds(schema, sourceSchema);
    Schema expectedSchema =
        new Schema(
            Lists.newArrayList(
                required(1, "a", Types.IntegerType.get()),
                Types.NestedField.required("c")
                    .withId(16)
                    .ofType(Types.IntegerType.get())
                    .withInitialDefault(Literal.of(23))
                    .withWriteDefault(Literal.of(34))
                    .build(),
                required(15, "B", Types.IntegerType.get())));

    assertThat(actualSchema.asStruct()).isEqualTo(expectedSchema.asStruct());
  }

  @Test
  public void testReassignOrRefreshIdsCaseInsensitive() {
    Schema schema =
        new Schema(
            Lists.newArrayList(
                required(1, "FIELD1", Types.IntegerType.get()),
                required(2, "FIELD2", Types.IntegerType.get())));
    Schema sourceSchema =
        new Schema(
            Lists.newArrayList(
                required(1, "field1", Types.IntegerType.get()),
                required(2, "field2", Types.IntegerType.get())));
    final Schema actualSchema = TypeUtil.reassignOrRefreshIds(schema, sourceSchema, false);
    final Schema expectedSchema =
        new Schema(
            Lists.newArrayList(
                required(1, "FIELD1", Types.IntegerType.get()),
                required(2, "FIELD2", Types.IntegerType.get())));
    assertThat(actualSchema.asStruct()).isEqualTo(expectedSchema.asStruct());
  }

  @Test
  public void testAssignIds() {
    Schema schema =
        new Schema(
            Lists.newArrayList(
                required(0, "a", Types.IntegerType.get()),
                Types.NestedField.required("c")
                    .withId(1)
                    .ofType(Types.IntegerType.get())
                    .withInitialDefault(Literal.of(23))
                    .withWriteDefault(Literal.of(34))
                    .build(),
                required(2, "B", Types.IntegerType.get())));

    Type actualSchema = TypeUtil.assignIds(schema.asStruct(), oldId -> oldId + 10);
    Schema expectedSchema =
        new Schema(
            Lists.newArrayList(
                required(10, "a", Types.IntegerType.get()),
                Types.NestedField.required("c")
                    .withId(11)
                    .ofType(Types.IntegerType.get())
                    .withInitialDefault(Literal.of(23))
                    .withWriteDefault(Literal.of(34))
                    .build(),
                required(12, "B", Types.IntegerType.get())));

    assertThat(actualSchema).isEqualTo(expectedSchema.asStruct());
  }

  @Test
  public void testAssignFreshIds() {
    Schema schema =
        new Schema(
            Lists.newArrayList(
                required(0, "a", Types.IntegerType.get()),
                Types.NestedField.required("c")
                    .withId(1)
                    .ofType(Types.IntegerType.get())
                    .withInitialDefault(Literal.of(23))
                    .withWriteDefault(Literal.of(34))
                    .build(),
                required(2, "B", Types.IntegerType.get())));

    Schema actualSchema = TypeUtil.assignFreshIds(schema, new AtomicInteger(10)::incrementAndGet);
    Schema expectedSchema =
        new Schema(
            Lists.newArrayList(
                required(11, "a", Types.IntegerType.get()),
                Types.NestedField.required("c")
                    .withId(12)
                    .ofType(Types.IntegerType.get())
                    .withInitialDefault(Literal.of(23))
                    .withWriteDefault(Literal.of(34))
                    .build(),
                required(13, "B", Types.IntegerType.get())));

    assertThat(actualSchema.asStruct()).isEqualTo(expectedSchema.asStruct());
  }

  @Test
  public void testReassignDoc() {
    Schema schema =
        new Schema(
            Lists.newArrayList(
                required(0, "a", Types.IntegerType.get()),
                Types.NestedField.required("c")
                    .withId(1)
                    .ofType(Types.IntegerType.get())
                    .withInitialDefault(Literal.of(23))
                    .withWriteDefault(Literal.of(34))
                    .build(),
                required(2, "B", Types.IntegerType.get())));

    Schema docSchema =
        new Schema(
            Lists.newArrayList(
                required(0, "a", Types.IntegerType.get(), "a_doc"),
                Types.NestedField.required("c")
                    .withId(1)
                    .ofType(Types.IntegerType.get())
                    .withDoc("c_doc")
                    .build(),
                required(2, "B", Types.IntegerType.get(), "b_doc")));

    Schema actualSchema = TypeUtil.reassignDoc(schema, docSchema);
    Schema expectedSchema =
        new Schema(
            Lists.newArrayList(
                required(0, "a", Types.IntegerType.get(), "a_doc"),
                Types.NestedField.required("c")
                    .withId(1)
                    .ofType(Types.IntegerType.get())
                    .withInitialDefault(Literal.of(23))
                    .withWriteDefault(Literal.of(34))
                    .withDoc("c_doc")
                    .build(),
                required(2, "B", Types.IntegerType.get(), "b_doc")));

    assertThat(actualSchema.asStruct()).isEqualTo(expectedSchema.asStruct());
  }

  private static Stream<Arguments> testTypes() {
    return Stream.of(
        Arguments.of(Types.UnknownType.get()),
        Arguments.of(Types.VariantType.get()),
        Arguments.of(Types.TimestampNanoType.withoutZone()),
        Arguments.of(Types.TimestampNanoType.withZone()),
        Arguments.of(Types.GeometryType.crs84()),
        Arguments.of(Types.GeometryType.of("srid:3857")),
        Arguments.of(Types.GeographyType.crs84()),
        Arguments.of(Types.GeographyType.of("srid:4269", EdgeAlgorithm.KARNEY)));
  }

  @ParameterizedTest
  @MethodSource("testTypes")
  public void testAssignIdsWithType(Type testType) {
    Types.StructType sourceType =
        Types.StructType.of(required(0, "id", IntegerType.get()), optional(1, "data", testType));
    Type expectedType =
        Types.StructType.of(required(10, "id", IntegerType.get()), optional(11, "data", testType));

    Type assignedType = TypeUtil.assignIds(sourceType, oldId -> oldId + 10);
    assertThat(assignedType).isEqualTo(expectedType);
  }

  @ParameterizedTest
  @MethodSource("testTypes")
  public void testAssignFreshIdsWithType(Type testType) {
    Schema schema = new Schema(required(0, "id", IntegerType.get()), optional(1, "data", testType));

    Schema assignedSchema = TypeUtil.assignFreshIds(schema, new AtomicInteger(10)::incrementAndGet);
    Schema expectedSchema =
        new Schema(required(11, "id", IntegerType.get()), optional(12, "data", testType));
    assertThat(assignedSchema.asStruct()).isEqualTo(expectedSchema.asStruct());
  }

  @ParameterizedTest
  @MethodSource("testTypes")
  public void testReassignIdsWithType(Type testType) {
    Schema schema = new Schema(required(0, "id", IntegerType.get()), optional(1, "data", testType));
    Schema sourceSchema =
        new Schema(required(1, "id", IntegerType.get()), optional(2, "data", testType));

    Schema reassignedSchema = TypeUtil.reassignIds(schema, sourceSchema);
    assertThat(reassignedSchema.asStruct()).isEqualTo(sourceSchema.asStruct());
  }

  @ParameterizedTest
  @MethodSource("testTypes")
  public void testIndexByIdWithType(Type testType) {
    Schema schema = new Schema(required(0, "id", IntegerType.get()), optional(1, "data", testType));

    Map<Integer, Types.NestedField> indexByIds = TypeUtil.indexById(schema.asStruct());
    assertThat(indexByIds.get(1).type()).isEqualTo(testType);
  }

  @ParameterizedTest
  @MethodSource("testTypes")
  public void testIndexNameByIdWithType(Type testType) {
    Schema schema = new Schema(required(0, "id", IntegerType.get()), optional(1, "data", testType));

    Map<Integer, String> indexNameByIds = TypeUtil.indexNameById(schema.asStruct());
    assertThat(indexNameByIds.get(1)).isEqualTo("data");
  }

  @ParameterizedTest
  @MethodSource("testTypes")
  public void testProjectWithType(Type testType) {
    Schema schema = new Schema(required(0, "id", IntegerType.get()), optional(1, "data", testType));

    Schema expectedSchema = new Schema(optional(1, "data", testType));
    Schema projectedSchema = TypeUtil.project(schema, Sets.newHashSet(1));
    assertThat(projectedSchema.asStruct()).isEqualTo(expectedSchema.asStruct());
  }

  @ParameterizedTest
  @MethodSource("testTypes")
  public void testGetProjectedIdsWithType(Type testType) {
    Schema schema = new Schema(required(0, "id", IntegerType.get()), optional(1, "data", testType));

    Set<Integer> projectedIds = TypeUtil.getProjectedIds(schema);
    assertThat(Set.of(0, 1)).isEqualTo(projectedIds);
  }

  @ParameterizedTest
  @MethodSource("testTypes")
  public void testReassignDocWithType(Type testType) {
    Schema schema = new Schema(required(0, "id", IntegerType.get()), optional(1, "data", testType));
    Schema docSourceSchema =
        new Schema(
            required(0, "id", IntegerType.get(), "id"), optional(1, "data", testType, "data"));

    Schema reassignedSchema = TypeUtil.reassignDoc(schema, docSourceSchema);
    assertThat(reassignedSchema.asStruct()).isEqualTo(docSourceSchema.asStruct());
  }

  @Test
  public void ancestorFieldsInEmptySchema() {
    assertThat(TypeUtil.ancestorFields(new Schema(), -1)).isEmpty();
    assertThat(TypeUtil.ancestorFields(new Schema(), 1)).isEmpty();
  }

  @Test
  public void ancestorFieldsInNonNestedSchema() {
    Schema schema =
        new Schema(
            required(0, "a", Types.IntegerType.get()), required(1, "A", Types.IntegerType.get()));

    assertThat(TypeUtil.ancestorFields(schema, 0)).isEmpty();
    assertThat(TypeUtil.ancestorFields(schema, 1)).isEmpty();
  }

  @Test
  public void ancestorFieldsInNestedSchema() {
    Types.NestedField innerPreferences =
        optional(
            8,
            "inner_preferences",
            Types.StructType.of(
                required(12, "feature3", Types.BooleanType.get()),
                optional(13, "feature4", Types.BooleanType.get())));
    Types.NestedField preferences =
        optional(
            3,
            "preferences",
            Types.StructType.of(
                required(6, "feature1", Types.BooleanType.get()),
                optional(7, "feature2", Types.BooleanType.get()),
                innerPreferences));

    // define locations map with all nested fields
    Types.StructType locationsKeyStruct =
        Types.StructType.of(
            required(20, "address", Types.StringType.get()),
            required(21, "city", Types.StringType.get()),
            required(22, "state", Types.StringType.get()),
            required(23, "zip", IntegerType.get()));
    Types.StructType locationsValueStruct =
        Types.StructType.of(
            required(14, "lat", Types.FloatType.get()),
            required(15, "long", Types.FloatType.get()));
    Types.NestedField locationsKey = required(9, "key", locationsKeyStruct);
    Types.NestedField locationsValue = required(10, "value", locationsValueStruct);
    Types.NestedField locations =
        required(
            4,
            "locations",
            Types.MapType.ofRequired(9, 10, locationsKeyStruct, locationsValueStruct));

    // define points list with all nested fields
    Types.StructType pointsStruct =
        Types.StructType.of(
            required(16, "x", Types.LongType.get()), required(17, "y", Types.LongType.get()));
    Types.NestedField pointsElement = optional(11, "element", pointsStruct);
    Types.NestedField points = optional(5, "points", Types.ListType.ofOptional(11, pointsStruct));

    Types.NestedField id = required(1, "id", IntegerType.get());
    Types.NestedField data = optional(2, "data", Types.StringType.get());
    Schema schema = new Schema(id, data, preferences, locations, points);

    // non-nested fields don't have parents
    assertThat(TypeUtil.ancestorFields(schema, id.fieldId())).isEmpty();
    assertThat(TypeUtil.ancestorFields(schema, data.fieldId())).isEmpty();

    // verify preferences struct and all of its nested fields (6-8, 12+13)
    assertThat(TypeUtil.ancestorFields(schema, preferences.fieldId())).isEmpty();
    assertThat(TypeUtil.ancestorFields(schema, 6)).containsExactly(preferences);
    assertThat(TypeUtil.ancestorFields(schema, 7)).containsExactly(preferences);
    assertThat(TypeUtil.ancestorFields(schema, innerPreferences.fieldId()))
        .containsExactly(preferences);
    assertThat(TypeUtil.ancestorFields(schema, 12)).containsExactly(innerPreferences, preferences);
    assertThat(TypeUtil.ancestorFields(schema, 13)).containsExactly(innerPreferences, preferences);

    // verify locations map and all of its nested fields (IDs 9+10, 20-23 and 14+15)
    assertThat(TypeUtil.ancestorFields(schema, locations.fieldId())).isEmpty();
    assertThat(TypeUtil.ancestorFields(schema, locationsKey.fieldId())).containsExactly(locations);
    assertThat(TypeUtil.ancestorFields(schema, 20)).containsExactly(locationsKey, locations);
    assertThat(TypeUtil.ancestorFields(schema, 21)).containsExactly(locationsKey, locations);
    assertThat(TypeUtil.ancestorFields(schema, 22)).containsExactly(locationsKey, locations);
    assertThat(TypeUtil.ancestorFields(schema, 23)).containsExactly(locationsKey, locations);
    assertThat(TypeUtil.ancestorFields(schema, locationsValue.fieldId()))
        .containsExactly(locations);
    assertThat(TypeUtil.ancestorFields(schema, 14)).containsExactly(locationsValue, locations);
    assertThat(TypeUtil.ancestorFields(schema, 15)).containsExactly(locationsValue, locations);

    // verify points list and all of its nested fields (IDs 11 and 16+17)
    assertThat(TypeUtil.ancestorFields(schema, points.fieldId())).isEmpty();
    assertThat(TypeUtil.ancestorFields(schema, pointsElement.fieldId())).containsExactly(points);
    assertThat(TypeUtil.ancestorFields(schema, 16)).containsExactly(pointsElement, points);
    assertThat(TypeUtil.ancestorFields(schema, 17)).containsExactly(pointsElement, points);
  }

  @Test
  public void testIndexStatsNames() {
    Schema schema =
        new Schema(
            required(1, "id", Types.LongType.get()),
            optional(2, "data", Types.StringType.get()),
            optional(
                3,
                "location",
                Types.StructType.of(
                    required(10, "lat", Types.FloatType.get()),
                    required(11, "long", Types.FloatType.get()))),
            optional(4, "values", Types.ListType.ofOptional(12, IntegerType.get())),
            optional(
                5,
                "points",
                Types.ListType.ofRequired(
                    13,
                    Types.StructType.of(
                        required(14, "x", IntegerType.get()),
                        required(15, "y", IntegerType.get())))),
            optional(
                6,
                "properties",
                Types.MapType.ofOptional(16, 17, Types.StringType.get(), Types.StringType.get())),
            optional(
                7,
                "addresses",
                Types.MapType.ofRequired(
                    18,
                    19,
                    Types.StringType.get(),
                    Types.StructType.of(
                        required(20, "street", Types.StringType.get()),
                        optional(21, "city", Types.StringType.get()),
                        optional(22, "state", Types.StringType.get()),
                        required(23, "zip", Types.IntegerType.get()),
                        required(24, "value", Types.IntegerType.get())))));

    Map<Integer, String> statsNameIndex = TypeUtil.indexStatsNames(schema.asStruct());

    assertThat(statsNameIndex)
        .containsEntry(1, "id")
        .containsEntry(2, "data")
        .containsEntry(3, "location")
        .containsEntry(10, "location_lat")
        .containsEntry(11, "location_long")
        .containsEntry(4, "values")
        .containsEntry(12, "values_element")
        .containsEntry(5, "points")
        .containsEntry(13, "points_element")
        .containsEntry(14, "points_x")
        .containsEntry(15, "points_y")
        .containsEntry(6, "properties")
        .containsEntry(16, "properties_key")
        .containsEntry(17, "properties_value")
        .containsEntry(7, "addresses")
        .containsEntry(18, "addresses_key")
        .containsEntry(19, "addresses_value")
        .containsEntry(20, "addresses_street")
        .containsEntry(21, "addresses_city")
        .containsEntry(22, "addresses_state")
        .containsEntry(23, "addresses_zip")
        .containsEntry(24, "addresses_value") // the leaf takes precedence
        .hasSize(22);
  }

  @Test
  public void testReplaceFieldTypes() {
    Types.StructType replacement = Types.StructType.of(required(10, "x", IntegerType.get()));
    Schema schema =
        new Schema(
            required(1, "id", IntegerType.get()),
            required(2, "s", Types.StructType.of(required(3, "a", Types.LongType.get()))));

    Schema result = TypeUtil.replaceFieldTypes(schema, ImmutableMap.of(2, replacement));

    assertThat(result.findField(1).type()).isEqualTo(IntegerType.get());
    assertThat(result.findField(2).type()).isEqualTo(replacement);
  }

  @Test
  public void testReplaceFieldTypesPrimitive() {
    Schema schema =
        new Schema(required(1, "id", IntegerType.get()), required(2, "count", IntegerType.get()));

    Schema result = TypeUtil.replaceFieldTypes(schema, ImmutableMap.of(2, Types.LongType.get()));

    assertThat(result.findField(1).type()).isEqualTo(IntegerType.get());
    assertThat(result.findField(2).type()).isEqualTo(Types.LongType.get());
  }

  @Test
  public void testReplaceFieldTypesListElement() {
    Schema schema =
        new Schema(required(1, "list", Types.ListType.ofRequired(2, Types.StructType.of())));

    Schema result =
        TypeUtil.replaceFieldTypes(
            schema, ImmutableMap.of(2, Types.StructType.of(required(3, "x", IntegerType.get()))));

    Types.ListType list = (Types.ListType) result.findField(1).type();
    assertThat(list.elementType().asStructType().field(3).name()).isEqualTo("x");
    assertThat(list.isElementRequired()).isTrue();
  }

  @Test
  public void testReplaceFieldTypesMapKeyAndValue() {
    Schema schema =
        new Schema(
            required(1, "m", Types.MapType.ofRequired(2, 3, IntegerType.get(), IntegerType.get())));

    Schema result =
        TypeUtil.replaceFieldTypes(
            schema, ImmutableMap.of(2, Types.LongType.get(), 3, Types.StringType.get()));

    Types.MapType map = (Types.MapType) result.findField(1).type();
    assertThat(map.keyType()).isEqualTo(Types.LongType.get());
    assertThat(map.valueType()).isEqualTo(Types.StringType.get());
    assertThat(map.isValueRequired()).isTrue();
  }

  @Test
  public void testReplaceFieldTypesPreservesDefaultsOnCompatibleType() {
    Types.NestedField field =
        Types.NestedField.optional("count")
            .withId(1)
            .ofType(IntegerType.get())
            .withInitialDefault(Literal.of(5))
            .withWriteDefault(Literal.of(7))
            .build();
    Schema schema = new Schema(field);

    Schema result = TypeUtil.replaceFieldTypes(schema, ImmutableMap.of(1, Types.LongType.get()));

    Types.NestedField replaced = result.findField(1);
    assertThat(replaced.type()).isEqualTo(Types.LongType.get());
    assertThat(replaced.initialDefault()).isEqualTo(5L);
    assertThat(replaced.writeDefault()).isEqualTo(7L);
  }

  @Test
  public void testReplaceFieldTypesRejectsDefaultIncompatibleWithType() {
    Types.NestedField field =
        Types.NestedField.optional("count")
            .withId(1)
            .ofType(IntegerType.get())
            .withInitialDefault(Literal.of(5))
            .build();
    Schema schema = new Schema(field);

    assertThatThrownBy(
            () -> TypeUtil.replaceFieldTypes(schema, ImmutableMap.of(1, Types.StringType.get())))
        .isInstanceOf(IllegalArgumentException.class)
        .hasMessageContaining("Cannot cast default value to string");
  }

  @Test
  public void testReplaceFieldTypesPreservesIdentifierField() {
    Schema schema =
        new Schema(
            Lists.newArrayList(
                required(1, "id", IntegerType.get()), required(2, "data", Types.StringType.get())),
            Sets.newHashSet(1));

    Schema result = TypeUtil.replaceFieldTypes(schema, ImmutableMap.of(1, Types.LongType.get()));

    assertThat(result.findField(1).type()).isEqualTo(Types.LongType.get());
    assertThat(result.identifierFieldIds()).containsExactly(1);
  }

  @Test
  public void testReplaceFieldTypesRejectsNonPrimitiveIdentifierField() {
    Schema schema =
        new Schema(
            Lists.newArrayList(
                required(1, "id", IntegerType.get()), required(2, "data", Types.StringType.get())),
            Sets.newHashSet(1));

    // an identifier field must be primitive, so replacing it with a struct is rejected
    assertThatThrownBy(
            () ->
                TypeUtil.replaceFieldTypes(
                    schema,
                    ImmutableMap.of(1, Types.StructType.of(required(3, "x", IntegerType.get())))))
        .isInstanceOf(IllegalArgumentException.class)
        .hasMessageContaining("not a primitive type field");
  }

  @Test
  public void testReplaceFieldTypesNoMatchReturnsSameSchema() {
    Schema schema = new Schema(required(1, "id", IntegerType.get()));
    Schema result = TypeUtil.replaceFieldTypes(schema, ImmutableMap.of(99, Types.LongType.get()));
    assertThat(result).isSameAs(schema);
  }
}
