/*
 * 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.spark.source;

import static org.apache.iceberg.TableProperties.DEFAULT_FILE_FORMAT;
import static org.apache.iceberg.TableProperties.FORMAT_VERSION;
import static org.apache.iceberg.TableProperties.ORC_VECTORIZATION_ENABLED;
import static org.apache.iceberg.TableProperties.PARQUET_BATCH_SIZE;
import static org.apache.iceberg.TableProperties.PARQUET_ROW_GROUP_SIZE_BYTES;
import static org.apache.iceberg.TableProperties.PARQUET_VECTORIZATION_ENABLED;
import static org.apache.spark.sql.functions.lit;

import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import org.apache.iceberg.FileFormat;
import org.apache.iceberg.HasTableOperations;
import org.apache.iceberg.MetadataColumns;
import org.apache.iceberg.PartitionSpec;
import org.apache.iceberg.PartitionSpecParser;
import org.apache.iceberg.Schema;
import org.apache.iceberg.Table;
import org.apache.iceberg.TableMetadata;
import org.apache.iceberg.TableOperations;
import org.apache.iceberg.exceptions.ValidationException;
import org.apache.iceberg.expressions.Expressions;
import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
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.Maps;
import org.apache.iceberg.spark.SparkTestBase;
import org.apache.iceberg.types.Types;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Encoders;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.catalyst.analysis.NoSuchTableException;
import org.assertj.core.api.Assertions;
import org.junit.After;
import org.junit.Assert;
import org.junit.Assume;
import org.junit.Before;
import org.junit.BeforeClass;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import org.junit.runner.RunWith;
import org.junit.runners.Parameterized;

@RunWith(Parameterized.class)
public class TestSparkMetadataColumns extends SparkTestBase {

  private static final String TABLE_NAME = "test_table";
  private static final Schema SCHEMA =
      new Schema(
          Types.NestedField.optional(1, "id", Types.LongType.get()),
          Types.NestedField.optional(2, "category", Types.StringType.get()),
          Types.NestedField.optional(3, "data", Types.StringType.get()));
  private static final PartitionSpec SPEC = PartitionSpec.unpartitioned();
  private static final PartitionSpec UNKNOWN_SPEC =
      PartitionSpecParser.fromJson(
          SCHEMA,
          "{ \"spec-id\": 1, \"fields\": [ { \"name\": \"id_zero\", \"transform\": \"zero\", \"source-id\": 1 } ] }");

  @Parameterized.Parameters(name = "fileFormat = {0}, vectorized = {1}, formatVersion = {2}")
  public static Object[][] parameters() {
    return new Object[][] {
      {FileFormat.PARQUET, false, 1},
      {FileFormat.PARQUET, true, 1},
      {FileFormat.PARQUET, false, 2},
      {FileFormat.PARQUET, true, 2},
      {FileFormat.AVRO, false, 1},
      {FileFormat.AVRO, false, 2},
      {FileFormat.ORC, false, 1},
      {FileFormat.ORC, true, 1},
      {FileFormat.ORC, false, 2},
      {FileFormat.ORC, true, 2},
    };
  }

  @Rule public TemporaryFolder temp = new TemporaryFolder();

  private final FileFormat fileFormat;
  private final boolean vectorized;
  private final int formatVersion;

  private Table table = null;

  public TestSparkMetadataColumns(FileFormat fileFormat, boolean vectorized, int formatVersion) {
    this.fileFormat = fileFormat;
    this.vectorized = vectorized;
    this.formatVersion = formatVersion;
  }

  @BeforeClass
  public static void setupSpark() {
    ImmutableMap<String, String> config =
        ImmutableMap.of(
            "type", "hive",
            "default-namespace", "default",
            "cache-enabled", "true");
    spark
        .conf()
        .set("spark.sql.catalog.spark_catalog", "org.apache.iceberg.spark.source.TestSparkCatalog");
    config.forEach(
        (key, value) -> spark.conf().set("spark.sql.catalog.spark_catalog." + key, value));
  }

  @Before
  public void setupTable() throws IOException {
    createAndInitTable();
  }

  @After
  public void dropTable() {
    TestTables.clearTables();
  }

  @Test
  public void testSpecAndPartitionMetadataColumns() {
    // TODO: support metadata structs in vectorized ORC reads
    Assume.assumeFalse(fileFormat == FileFormat.ORC && vectorized);

    sql("INSERT INTO TABLE %s VALUES (1, 'a1', 'b1')", TABLE_NAME);

    table.refresh();
    table.updateSpec().addField("data").commit();
    sql("INSERT INTO TABLE %s VALUES (1, 'a1', 'b1')", TABLE_NAME);

    table.refresh();
    table.updateSpec().addField(Expressions.bucket("category", 8)).commit();
    sql("INSERT INTO TABLE %s VALUES (1, 'a1', 'b1')", TABLE_NAME);

    table.refresh();
    table.updateSpec().removeField("data").commit();
    sql("INSERT INTO TABLE %s VALUES (1, 'a1', 'b1')", TABLE_NAME);

    table.refresh();
    table.updateSpec().renameField("category_bucket_8", "category_bucket_8_another_name").commit();

    List<Object[]> expected =
        ImmutableList.of(
            row(0, row(null, null)),
            row(1, row("b1", null)),
            row(2, row("b1", 2)),
            row(3, row(null, 2)));
    assertEquals(
        "Rows must match",
        expected,
        sql("SELECT _spec_id, _partition FROM %s ORDER BY _spec_id", TABLE_NAME));
  }

  @Test
  public void testPositionMetadataColumnWithMultipleRowGroups() throws NoSuchTableException {
    Assume.assumeTrue(fileFormat == FileFormat.PARQUET);

    table.updateProperties().set(PARQUET_ROW_GROUP_SIZE_BYTES, "100").commit();

    List<Long> ids = Lists.newArrayList();
    for (long id = 0L; id < 200L; id++) {
      ids.add(id);
    }
    Dataset<Row> df =
        spark
            .createDataset(ids, Encoders.LONG())
            .withColumnRenamed("value", "id")
            .withColumn("category", lit("hr"))
            .withColumn("data", lit("ABCDEF"));
    df.coalesce(1).writeTo(TABLE_NAME).append();

    Assert.assertEquals(200, spark.table(TABLE_NAME).count());

    List<Object[]> expectedRows = ids.stream().map(this::row).collect(Collectors.toList());
    assertEquals("Rows must match", expectedRows, sql("SELECT _pos FROM %s", TABLE_NAME));
  }

  @Test
  public void testPositionMetadataColumnWithMultipleBatches() throws NoSuchTableException {
    Assume.assumeTrue(fileFormat == FileFormat.PARQUET);

    table.updateProperties().set(PARQUET_BATCH_SIZE, "1000").commit();

    List<Long> ids = Lists.newArrayList();
    for (long id = 0L; id < 7500L; id++) {
      ids.add(id);
    }
    Dataset<Row> df =
        spark
            .createDataset(ids, Encoders.LONG())
            .withColumnRenamed("value", "id")
            .withColumn("category", lit("hr"))
            .withColumn("data", lit("ABCDEF"));
    df.coalesce(1).writeTo(TABLE_NAME).append();

    Assert.assertEquals(7500, spark.table(TABLE_NAME).count());

    List<Object[]> expectedRows = ids.stream().map(this::row).collect(Collectors.toList());
    assertEquals("Rows must match", expectedRows, sql("SELECT _pos FROM %s", TABLE_NAME));
  }

  @Test
  public void testPartitionMetadataColumnWithUnknownTransforms() {
    // replace the table spec to include an unknown transform
    TableOperations ops = ((HasTableOperations) table).operations();
    TableMetadata base = ops.current();
    ops.commit(base, base.updatePartitionSpec(UNKNOWN_SPEC));

    Assertions.assertThatThrownBy(() -> sql("SELECT _partition FROM %s", TABLE_NAME))
        .isInstanceOf(ValidationException.class)
        .hasMessage("Cannot build table partition type, unknown transforms: [zero]");
  }

  @Test
  public void testConflictingColumns() {
    table
        .updateSchema()
        .addColumn(MetadataColumns.SPEC_ID.name(), Types.IntegerType.get())
        .addColumn(MetadataColumns.FILE_PATH.name(), Types.StringType.get())
        .commit();

    sql("INSERT INTO TABLE %s VALUES (1, 'a1', 'b1', -1, 'path/to/file')", TABLE_NAME);

    assertEquals(
        "Rows must match",
        ImmutableList.of(row(1L, "a1")),
        sql("SELECT id, category FROM %s", TABLE_NAME));

    Assertions.assertThatThrownBy(() -> sql("SELECT * FROM %s", TABLE_NAME))
        .isInstanceOf(ValidationException.class)
        .hasMessageStartingWith(
            "Table column names conflict with names reserved for Iceberg metadata columns: [_spec_id, _file].");

    table.refresh();

    table
        .updateSchema()
        .renameColumn(MetadataColumns.SPEC_ID.name(), "_renamed" + MetadataColumns.SPEC_ID.name())
        .renameColumn(
            MetadataColumns.FILE_PATH.name(), "_renamed" + MetadataColumns.FILE_PATH.name())
        .commit();

    assertEquals(
        "Rows must match",
        ImmutableList.of(row(0, null, -1)),
        sql("SELECT _spec_id, _partition, _renamed_spec_id FROM %s", TABLE_NAME));
  }

  private void createAndInitTable() throws IOException {
    Map<String, String> properties = Maps.newHashMap();
    properties.put(FORMAT_VERSION, String.valueOf(formatVersion));
    properties.put(DEFAULT_FILE_FORMAT, fileFormat.name());

    switch (fileFormat) {
      case PARQUET:
        properties.put(PARQUET_VECTORIZATION_ENABLED, String.valueOf(vectorized));
        break;
      case ORC:
        properties.put(ORC_VECTORIZATION_ENABLED, String.valueOf(vectorized));
        break;
      default:
        Preconditions.checkState(
            !vectorized, "File format %s does not support vectorized reads", fileFormat);
    }

    this.table = TestTables.create(temp.newFolder(), TABLE_NAME, SCHEMA, SPEC, properties);
  }
}
