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

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

import java.io.File;
import java.io.IOException;
import java.util.List;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.Path;
import org.apache.iceberg.DataFile;
import org.apache.iceberg.DeleteFile;
import org.apache.iceberg.FileScanTask;
import org.apache.iceberg.Files;
import org.apache.iceberg.SnapshotChanges;
import org.apache.iceberg.Table;
import org.apache.iceberg.data.FileHelpers;
import org.apache.iceberg.io.CloseableIterable;
import org.apache.iceberg.relocated.com.google.common.collect.Iterables;
import org.apache.iceberg.relocated.com.google.common.collect.Lists;
import org.apache.iceberg.spark.Spark3Util;
import org.apache.iceberg.spark.SparkCatalog;
import org.apache.iceberg.spark.TestBase;
import org.apache.iceberg.util.CharSequenceSet;
import org.apache.iceberg.util.Pair;
import org.apache.parquet.hadoop.ParquetFileReader;
import org.apache.parquet.hadoop.util.HadoopInputFile;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.catalyst.analysis.NoSuchTableException;
import org.apache.spark.sql.catalyst.parser.ParseException;
import org.apache.spark.types.variant.Variant;
import org.apache.spark.unsafe.types.VariantVal;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;

public class TestSparkVariantRead extends TestBase {

  private static final String CATALOG = "local";
  private static final String TABLE = CATALOG + ".default.var";

  @BeforeAll
  public static void setupCatalog() {
    // Use a Hadoop catalog to avoid Hive schema conversion (Hive doesn't support VARIANT yet)
    spark.conf().set("spark.sql.catalog." + CATALOG, SparkCatalog.class.getName());
    spark.conf().set("spark.sql.catalog." + CATALOG + ".type", "hadoop");
    spark.conf().set("spark.sql.catalog." + CATALOG + ".default-namespace", "default");
    spark.conf().set("spark.sql.catalog." + CATALOG + ".cache-enabled", "false");
    // point warehouse to a temp directory
    String temp = System.getProperty("java.io.tmpdir") + "/iceberg_spark_variant_warehouse";
    spark.conf().set("spark.sql.catalog." + CATALOG + ".warehouse", temp);
  }

  @BeforeEach
  public void setupTable() {
    sql("DROP TABLE IF EXISTS %s", TABLE);
    sql(
        "CREATE TABLE %s (id BIGINT, v1 VARIANT, v2 VARIANT) USING iceberg "
            + "TBLPROPERTIES ('format-version'='3')",
        TABLE);

    String v1r1 = "{\"a\":1}";
    String v2r1 = "{\"x\":10}";
    String v1r2 = "{\"b\":2}";
    String v2r2 = "{\"y\":20}";

    sql("INSERT INTO %s SELECT 1, parse_json('%s'), parse_json('%s')", TABLE, v1r1, v2r1);
    sql("INSERT INTO %s SELECT 2, parse_json('%s'), parse_json('%s')", TABLE, v1r2, v2r2);
  }

  @AfterEach
  public void cleanup() {
    sql("DROP TABLE IF EXISTS %s", TABLE);
  }

  @ParameterizedTest
  @ValueSource(booleans = {false, true})
  public void testVariantColumnProjection_singleVariant(boolean vectorized) {
    setVectorization(vectorized);
    Dataset<Row> df = spark.table(TABLE).select("id", "v1").orderBy("id");
    assertThat(df.schema().fieldNames()).containsExactly("id", "v1");
    assertThat(df.count()).isEqualTo(2);

    List<Row> directRows = df.collectAsList();
    Object v1row1 = directRows.get(0).get(1);
    Object v1row2 = directRows.get(1).get(1);
    assertThat(v1row1).isInstanceOf(VariantVal.class);
    assertThat(v1row2).isInstanceOf(VariantVal.class);
    VariantVal r1 = (VariantVal) v1row1;
    VariantVal r2 = (VariantVal) v1row2;
    Variant vv1 = new Variant(r1.getValue(), r1.getMetadata());
    Variant vv2 = new Variant(r2.getValue(), r2.getMetadata());

    // row 1 has {"a":1}
    Variant fieldA = vv1.getFieldByKey("a");
    assertThat(fieldA).isNotNull();
    assertThat(fieldA.getLong()).isEqualTo(1L);

    // row 2 has {"b":2}
    Variant fieldB = vv2.getFieldByKey("b");
    assertThat(fieldB).isNotNull();
    assertThat(fieldB.getLong()).isEqualTo(2L);
  }

  @ParameterizedTest
  @ValueSource(booleans = {false, true})
  public void testVariantColumnProjectionNoVariant(boolean vectorized) {
    setVectorization(vectorized);
    Dataset<Row> df = spark.table(TABLE).select("id");
    assertThat(df.schema().fieldNames()).containsExactly("id");
    assertThat(df.count()).isEqualTo(2);
    assertThat(df.collectAsList()).extracting(r -> r.getLong(0)).containsExactlyInAnyOrder(1L, 2L);
  }

  @ParameterizedTest
  @ValueSource(booleans = {false, true})
  public void testFilterOnVariantColumnOnWholeValue(boolean vectorized) {
    setVectorization(vectorized);
    sql("INSERT INTO %s SELECT 3, NULL, NULL", TABLE);

    Dataset<Row> nullDf = spark.table(TABLE).where("v1 IS NULL").select("id");
    assertThat(nullDf.collectAsList()).extracting(r -> r.getLong(0)).containsExactly(3L);

    Dataset<Row> notNullDf = spark.table(TABLE).where("v1 IS NOT NULL").select("id");
    assertThat(notNullDf.collectAsList())
        .extracting(r -> r.getLong(0))
        .containsExactlyInAnyOrder(1L, 2L);

    // verify variant contents for non-null rows
    Dataset<Row> notNullVals =
        spark
            .table(TABLE)
            .where("v1 IS NOT NULL")
            .selectExpr("id", "to_json(v1) as v1_json")
            .orderBy("id");
    List<Row> nn = notNullVals.collectAsList();
    assertThat(nn).hasSize(2);
    assertThat(nn.get(0).getLong(0)).isEqualTo(1L);
    assertThat(nn.get(0).getString(1)).isEqualTo("{\"a\":1}");
    assertThat(nn.get(1).getLong(0)).isEqualTo(2L);
    assertThat(nn.get(1).getString(1)).isEqualTo("{\"b\":2}");
  }

  @ParameterizedTest
  @ValueSource(booleans = {false, true})
  public void testVariantNullValueProjection(boolean vectorized) {
    setVectorization(vectorized);

    // insert a row with NULL variant values
    sql("INSERT INTO %s SELECT 10, NULL, NULL", TABLE);

    // select id and variant; ensure the variant value is null
    Dataset<Row> df = spark.table(TABLE).where("id = 10").select("id", "v1");
    List<Row> rows = df.collectAsList();
    assertThat(rows).hasSize(1);
    Row row = rows.get(0);
    assertThat(row.getLong(0)).isEqualTo(10L);
    assertThat(row.isNullAt(1)).isTrue();
  }

  @ParameterizedTest
  @ValueSource(booleans = {false, true})
  public void testVariantReadAfterDelete(boolean vectorized)
      throws IOException, NoSuchTableException, ParseException {
    String deleteTable = CATALOG + ".default.var_delete";

    sql("DROP TABLE IF EXISTS %s", deleteTable);
    sql(
        "CREATE TABLE %s (id BIGINT, v1 VARIANT) USING iceberg "
            + "TBLPROPERTIES ('format-version'='3')",
        deleteTable);
    setVectorization(deleteTable, vectorized);

    spark
        .sql(
            "SELECT 1L AS id, parse_json('{\"a\":1}') AS v1 "
                + "UNION ALL SELECT 2L, parse_json('{\"b\":2}')")
        .coalesce(1)
        .writeTo(deleteTable)
        .append();

    Table table = Spark3Util.loadIcebergTable(spark, deleteTable);
    DataFile dataFile =
        Iterables.getOnlyElement(SnapshotChanges.builderFor(table).build().addedDataFiles());

    Pair<DeleteFile, CharSequenceSet> deletes =
        FileHelpers.writeDeleteFile(
            table,
            Files.localOutput(File.createTempFile("dv-", ".puffin")),
            null,
            Lists.newArrayList(Pair.of(dataFile.location(), 0L)),
            3);

    table
        .newRowDelta()
        .addDeletes(deletes.first())
        .validateDataFilesExist(deletes.second())
        .commit();

    sql("REFRESH TABLE %s", deleteTable);

    Dataset<Row> df = spark.table(deleteTable).select("id", "v1").orderBy("id");
    List<Row> rows = df.collectAsList();

    assertThat(rows).hasSize(1);
    assertThat(rows.get(0).getLong(0)).isEqualTo(2L);

    Variant v1 =
        new Variant(
            ((VariantVal) rows.get(0).get(1)).getValue(),
            ((VariantVal) rows.get(0).get(1)).getMetadata());
    assertThat(v1.getFieldByKey("b").getLong()).isEqualTo(2L);

    sql("DROP TABLE IF EXISTS %s", deleteTable);
  }

  @ParameterizedTest
  @ValueSource(booleans = {false, true})
  public void testNestedStructVariant(boolean vectorized) {

    String structTable = CATALOG + ".default.var_struct";
    sql("DROP TABLE IF EXISTS %s", structTable);
    sql(
        "CREATE TABLE %s (id BIGINT, s STRUCT<v: VARIANT>) USING iceberg "
            + "TBLPROPERTIES ('format-version'='3')",
        structTable);
    setVectorization(structTable, vectorized);

    String j1 = "{\"a\":1}";
    String j2 = "{\"b\":2}";
    sql("INSERT INTO %s SELECT 1, named_struct('v', parse_json('%s'))", structTable, j1);
    sql("INSERT INTO %s SELECT 2, named_struct('v', parse_json('%s'))", structTable, j2);

    Dataset<Row> df = spark.table(structTable).selectExpr("id", "s.v AS v").orderBy("id");
    java.util.List<Row> rows = df.collectAsList();
    assertThat(rows.get(0).getLong(0)).isEqualTo(1L);
    Object sv1 = rows.get(0).get(1);
    assertThat(sv1).isInstanceOf(VariantVal.class);
    Variant sv1Var = new Variant(((VariantVal) sv1).getValue(), ((VariantVal) sv1).getMetadata());
    assertThat(sv1Var.getFieldByKey("a").getLong()).isEqualTo(1L);

    assertThat(rows.get(1).getLong(0)).isEqualTo(2L);
    Object sv2 = rows.get(1).get(1);
    assertThat(sv2).isInstanceOf(VariantVal.class);
    Variant sv2Var = new Variant(((VariantVal) sv2).getValue(), ((VariantVal) sv2).getMetadata());
    assertThat(sv2Var.getFieldByKey("b").getLong()).isEqualTo(2L);

    sql("DROP TABLE IF EXISTS %s", structTable);
  }

  @ParameterizedTest
  @ValueSource(booleans = {false, true})
  public void testNestedArrayVariant(boolean vectorized) {

    String arrayTable = CATALOG + ".default.var_array";
    sql("DROP TABLE IF EXISTS %s", arrayTable);
    sql(
        "CREATE TABLE %s (id BIGINT, arr ARRAY<VARIANT>) USING iceberg "
            + "TBLPROPERTIES ('format-version'='3')",
        arrayTable);
    setVectorization(arrayTable, vectorized);

    String a1 = "{\"a\":1}";
    String a2 = "{\"x\":10}";
    String b1 = "{\"b\":2}";
    String b2 = "{\"y\":20}";
    sql("INSERT INTO %s SELECT 1, array(parse_json('%s'), parse_json('%s'))", arrayTable, a1, a2);
    sql("INSERT INTO %s SELECT 2, array(parse_json('%s'), parse_json('%s'))", arrayTable, b1, b2);

    Dataset<Row> df =
        spark.table(arrayTable).selectExpr("id", "arr[0] as e0", "arr[1] as e1").orderBy("id");
    java.util.List<Row> rows = df.collectAsList();
    assertThat(rows.get(0).getLong(0)).isEqualTo(1L);
    Variant e0r1 =
        new Variant(
            ((VariantVal) rows.get(0).get(1)).getValue(),
            ((VariantVal) rows.get(0).get(1)).getMetadata());
    Variant e1r1 =
        new Variant(
            ((VariantVal) rows.get(0).get(2)).getValue(),
            ((VariantVal) rows.get(0).get(2)).getMetadata());
    assertThat(e0r1.getFieldByKey("a").getLong()).isEqualTo(1L);
    assertThat(e1r1.getFieldByKey("x").getLong()).isEqualTo(10L);
    assertThat(rows.get(1).getLong(0)).isEqualTo(2L);
    Variant e0r2 =
        new Variant(
            ((VariantVal) rows.get(1).get(1)).getValue(),
            ((VariantVal) rows.get(1).get(1)).getMetadata());
    Variant e1r2 =
        new Variant(
            ((VariantVal) rows.get(1).get(2)).getValue(),
            ((VariantVal) rows.get(1).get(2)).getMetadata());
    assertThat(e0r2.getFieldByKey("b").getLong()).isEqualTo(2L);
    assertThat(e1r2.getFieldByKey("y").getLong()).isEqualTo(20L);

    sql("DROP TABLE IF EXISTS %s", arrayTable);
  }

  @ParameterizedTest
  @ValueSource(booleans = {false, true})
  public void testNestedMapVariant(boolean vectorized) {

    String mapTable = CATALOG + ".default.var_map";
    sql("DROP TABLE IF EXISTS %s", mapTable);
    sql(
        "CREATE TABLE %s (id BIGINT, m MAP<STRING, VARIANT>) USING iceberg "
            + "TBLPROPERTIES ('format-version'='3')",
        mapTable);
    setVectorization(mapTable, vectorized);

    String k1a = "{\"a\":1}";
    String k2x = "{\"x\":10}";
    String k1b = "{\"b\":2}";
    String k2y = "{\"y\":20}";
    sql(
        "INSERT INTO %s SELECT 1, map('k1', parse_json('%s'), 'k2', parse_json('%s'))",
        mapTable, k1a, k2x);
    sql(
        "INSERT INTO %s SELECT 2, map('k1', parse_json('%s'), 'k2', parse_json('%s'))",
        mapTable, k1b, k2y);

    Dataset<Row> df =
        spark
            .table(mapTable)
            .selectExpr("id", "element_at(m, 'k1') as k1", "element_at(m, 'k2') as k2")
            .orderBy("id");
    java.util.List<Row> rows = df.collectAsList();
    assertThat(rows.get(0).getLong(0)).isEqualTo(1L);
    Variant k1r1 =
        new Variant(
            ((VariantVal) rows.get(0).get(1)).getValue(),
            ((VariantVal) rows.get(0).get(1)).getMetadata());
    Variant k2r1 =
        new Variant(
            ((VariantVal) rows.get(0).get(2)).getValue(),
            ((VariantVal) rows.get(0).get(2)).getMetadata());
    assertThat(k1r1.getFieldByKey("a").getLong()).isEqualTo(1L);
    assertThat(k2r1.getFieldByKey("x").getLong()).isEqualTo(10L);
    assertThat(rows.get(1).getLong(0)).isEqualTo(2L);
    Variant k1r2 =
        new Variant(
            ((VariantVal) rows.get(1).get(1)).getValue(),
            ((VariantVal) rows.get(1).get(1)).getMetadata());
    Variant k2r2 =
        new Variant(
            ((VariantVal) rows.get(1).get(2)).getValue(),
            ((VariantVal) rows.get(1).get(2)).getMetadata());
    assertThat(k1r2.getFieldByKey("b").getLong()).isEqualTo(2L);
    assertThat(k2r2.getFieldByKey("y").getLong()).isEqualTo(20L);

    sql("DROP TABLE IF EXISTS %s", mapTable);
  }

  @ParameterizedTest
  @ValueSource(booleans = {false, true})
  public void testMergeIntoWithVariant(boolean vectorized) {
    String mergeTable = CATALOG + ".default.var_merge";
    sql("DROP TABLE IF EXISTS %s", mergeTable);
    sql(
        "CREATE TABLE %s (id BIGINT, data VARIANT) USING iceberg "
            + "TBLPROPERTIES ('format-version'='3')",
        mergeTable);
    setVectorization(mergeTable, vectorized);

    sql(
        "INSERT INTO %s VALUES "
            + "(1, parse_json('{\"name\":\"alice\",\"age\":30}')), "
            + "(2, parse_json('{\"name\":\"bob\",\"age\":25}'))",
        mergeTable);

    sql(
        "MERGE INTO %s AS target "
            + "USING (SELECT 1 AS id, parse_json('{\"name\":\"alice\",\"age\":31}') AS data) AS source "
            + "ON target.id = source.id "
            + "WHEN MATCHED THEN UPDATE SET target.data = source.data "
            + "WHEN NOT MATCHED THEN INSERT *",
        mergeTable);

    List<Row> rows = spark.table(mergeTable).select("id", "data").orderBy("id").collectAsList();

    assertThat(rows).hasSize(2);
    assertThat(rows.get(0).getLong(0)).isEqualTo(1L);
    Variant v1 =
        new Variant(
            ((VariantVal) rows.get(0).get(1)).getValue(),
            ((VariantVal) rows.get(0).get(1)).getMetadata());
    assertThat(v1.getFieldByKey("name").getString()).describedAs("v1.name").isEqualTo("alice");
    assertThat(v1.getFieldByKey("age").getLong()).describedAs("v1.age").isEqualTo(31L);

    assertThat(rows.get(1).getLong(0)).isEqualTo(2L);
    Variant v2 =
        new Variant(
            ((VariantVal) rows.get(1).get(1)).getValue(),
            ((VariantVal) rows.get(1).get(1)).getMetadata());
    assertThat(v2.getFieldByKey("name").getString()).describedAs("v2.name").isEqualTo("bob");
    assertThat(v2.getFieldByKey("age").getLong()).describedAs("v2.age").isEqualTo(25L);

    sql("DROP TABLE IF EXISTS %s", mergeTable);
  }

  @ParameterizedTest
  @ValueSource(booleans = {false, true})
  public void testReadShreddedAfterPropertyToggled(boolean vectorized)
      throws IOException, NoSuchTableException, ParseException {
    String toggleTable = CATALOG + ".default.var_toggle";
    sql("DROP TABLE IF EXISTS %s", toggleTable);
    sql(
        "CREATE TABLE %s (id BIGINT, v VARIANT) USING iceberg "
            + "TBLPROPERTIES ('format-version'='3', 'write.parquet.shred-variants'='true')",
        toggleTable);

    spark.conf().set("spark.sql.iceberg.shred-variants", "true");
    try {
      sql(
          "INSERT INTO %s VALUES "
              + "(1, parse_json('{\"name\":\"alice\",\"age\":30}')), "
              + "(2, parse_json('{\"name\":\"bob\",\"age\":25}'))",
          toggleTable);
    } finally {
      spark.conf().unset("spark.sql.iceberg.shred-variants");
    }

    Table table = Spark3Util.loadIcebergTable(spark, toggleTable);
    assertHasTypedValueSubtree(table);

    sql("ALTER TABLE %s SET TBLPROPERTIES ('write.parquet.shred-variants'='false')", toggleTable);
    setVectorization(toggleTable, vectorized);

    List<Row> rows = spark.table(toggleTable).select("id", "v").orderBy("id").collectAsList();
    assertThat(rows).hasSize(2);
    Variant v1 =
        new Variant(
            ((VariantVal) rows.get(0).get(1)).getValue(),
            ((VariantVal) rows.get(0).get(1)).getMetadata());
    assertThat(v1.getFieldByKey("name").getString()).isEqualTo("alice");
    assertThat(v1.getFieldByKey("age").getLong()).isEqualTo(30L);
    Variant v2 =
        new Variant(
            ((VariantVal) rows.get(1).get(1)).getValue(),
            ((VariantVal) rows.get(1).get(1)).getMetadata());
    assertThat(v2.getFieldByKey("name").getString()).isEqualTo("bob");
    assertThat(v2.getFieldByKey("age").getLong()).isEqualTo(25L);

    sql("DROP TABLE IF EXISTS %s", toggleTable);
  }

  @ParameterizedTest
  @ValueSource(strings = {"none", "counts"})
  public void testReadShreddedWithMetricsDisabled(String metricsMode)
      throws IOException, NoSuchTableException, ParseException {
    String noStatsTable = CATALOG + ".default.var_no_stats";
    sql("DROP TABLE IF EXISTS %s", noStatsTable);
    sql(
        "CREATE TABLE %s (id BIGINT, v VARIANT) USING iceberg "
            + "TBLPROPERTIES ("
            + "'format-version'='3', "
            + "'write.parquet.shred-variants'='true', "
            + "'write.metadata.metrics.default'='%s')",
        noStatsTable, metricsMode);

    spark.conf().set("spark.sql.iceberg.shred-variants", "true");
    try {
      sql(
          "INSERT INTO %s VALUES "
              + "(1, parse_json('{\"name\":\"alice\",\"age\":30}')), "
              + "(2, parse_json('{\"name\":\"bob\",\"age\":25}'))",
          noStatsTable);
    } finally {
      spark.conf().unset("spark.sql.iceberg.shred-variants");
    }

    Table table = Spark3Util.loadIcebergTable(spark, noStatsTable);
    assertHasTypedValueSubtree(table);
    setVectorization(noStatsTable, true);

    List<Row> rows = spark.table(noStatsTable).select("id", "v").orderBy("id").collectAsList();
    assertThat(rows).hasSize(2);
    Variant v1 =
        new Variant(
            ((VariantVal) rows.get(0).get(1)).getValue(),
            ((VariantVal) rows.get(0).get(1)).getMetadata());
    assertThat(v1.getFieldByKey("name").getString()).isEqualTo("alice");
    Variant v2 =
        new Variant(
            ((VariantVal) rows.get(1).get(1)).getValue(),
            ((VariantVal) rows.get(1).get(1)).getMetadata());
    assertThat(v2.getFieldByKey("name").getString()).isEqualTo("bob");

    sql("DROP TABLE IF EXISTS %s", noStatsTable);
  }

  private void setVectorization(boolean on) {
    sql(
        "ALTER TABLE %s SET TBLPROPERTIES ('read.parquet.vectorization.enabled'='%s')",
        TABLE, Boolean.toString(on));
  }

  private void setVectorization(String table, boolean on) {
    sql(
        "ALTER TABLE %s SET TBLPROPERTIES ('read.parquet.vectorization.enabled'='%s')",
        table, Boolean.toString(on));
  }

  private static void assertHasTypedValueSubtree(Table table) throws IOException {
    try (CloseableIterable<FileScanTask> tasks = table.newScan().planFiles()) {
      assertThat(tasks).isNotEmpty();
      for (FileScanTask task : tasks) {
        HadoopInputFile inputFile =
            HadoopInputFile.fromPath(new Path(task.file().location()), new Configuration());
        try (ParquetFileReader reader = ParquetFileReader.open(inputFile)) {
          assertThat(containsTypedValue(reader.getFileMetaData().getSchema()))
              .as("Expected variant column to be shredded with a typed_value subtree")
              .isTrue();
        }
      }
    }
  }

  private static boolean containsTypedValue(org.apache.parquet.schema.Type type) {
    if (type.isPrimitive()) {
      return false;
    }
    for (org.apache.parquet.schema.Type child : type.asGroupType().getFields()) {
      if (child.getName().equals("typed_value") || containsTypedValue(child)) {
        return true;
      }
    }
    return false;
  }
}
