/*
 * 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 static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.assertj.core.api.Assumptions.assumeThat;

import java.util.List;
import java.util.Map;
import org.apache.iceberg.Parameter;
import org.apache.iceberg.ParameterizedTestExtension;
import org.apache.iceberg.Parameters;
import org.apache.iceberg.RowLevelOperationMode;
import org.apache.iceberg.SnapshotSummary;
import org.apache.iceberg.Table;
import org.apache.iceberg.TableProperties;
import org.apache.iceberg.TestHelpers;
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.spark.CatalogTestBase;
import org.apache.iceberg.spark.source.SimpleRecord;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.catalyst.analysis.NoSuchTableException;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.TestTemplate;
import org.junit.jupiter.api.extension.ExtendWith;

@ExtendWith(ParameterizedTestExtension.class)
public class TestDeleteFrom extends CatalogTestBase {
  @Parameter(index = 3)
  private int formatVersion;

  @Parameters(name = "catalogName = {0}, implementation = {1}, config = {2}, formatVersion = {3}")
  protected static Object[][] parameters() {
    List<Object[]> parameters = Lists.newArrayList();
    for (Object[] catalogParams : CatalogTestBase.parameters()) {
      for (int version : TestHelpers.V2_AND_ABOVE) {
        parameters.add(
            new Object[] {catalogParams[0], catalogParams[1], catalogParams[2], version});
      }
    }

    return parameters.toArray(new Object[0][]);
  }

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

  private String tableProperties() {
    return tableProperties(ImmutableMap.of());
  }

  private String tableProperties(Map<String, String> additionalProperties) {
    ImmutableMap.Builder<String, String> builder = ImmutableMap.builder();
    builder.putAll(additionalProperties);
    builder.put(TableProperties.FORMAT_VERSION, String.valueOf(formatVersion));
    return String.format("TBLPROPERTIES (%s)", tablePropsAsString(builder.buildKeepingLast()));
  }

  @TestTemplate
  public void testDeleteFromUnpartitionedTable() throws NoSuchTableException {
    sql("CREATE TABLE %s (id bigint, data string) USING iceberg %s", tableName, tableProperties());

    List<SimpleRecord> records =
        Lists.newArrayList(
            new SimpleRecord(1, "a"), new SimpleRecord(2, "b"), new SimpleRecord(3, "c"));
    Dataset<Row> df = spark.createDataFrame(records, SimpleRecord.class);
    df.coalesce(1).writeTo(tableName).append();

    assertEquals(
        "Should have expected rows",
        ImmutableList.of(row(1L, "a"), row(2L, "b"), row(3L, "c")),
        sql("SELECT * FROM %s ORDER BY id", tableName));

    sql("DELETE FROM %s WHERE id < 2", tableName);

    assertEquals(
        "Should have no rows after successful delete",
        ImmutableList.of(row(2L, "b"), row(3L, "c")),
        sql("SELECT * FROM %s ORDER BY id", tableName));

    sql("DELETE FROM %s WHERE id < 4", tableName);

    assertEquals(
        "Should have no rows after successful delete",
        ImmutableList.of(),
        sql("SELECT * FROM %s ORDER BY id", tableName));
  }

  @TestTemplate
  public void testDeleteFromTableAtSnapshot() throws NoSuchTableException {
    sql("CREATE TABLE %s (id bigint, data string) USING iceberg %s", tableName, tableProperties());

    List<SimpleRecord> records =
        Lists.newArrayList(
            new SimpleRecord(1, "a"), new SimpleRecord(2, "b"), new SimpleRecord(3, "c"));
    Dataset<Row> df = spark.createDataFrame(records, SimpleRecord.class);
    df.coalesce(1).writeTo(tableName).append();

    long snapshotId = validationCatalog.loadTable(tableIdent).currentSnapshot().snapshotId();
    String prefix = "snapshot_id_";
    assertThatThrownBy(() -> sql("DELETE FROM %s.%s WHERE id < 4", tableName, prefix + snapshotId))
        .isInstanceOf(IllegalArgumentException.class)
        .hasMessageStartingWith("Cannot modify table with time travel");
  }

  @TestTemplate
  public void testDeleteFromPartitionedTable() throws NoSuchTableException {
    sql(
        "CREATE TABLE %s (id bigint, data string) "
            + "USING iceberg "
            + "PARTITIONED BY (truncate(id, 2)) "
            + "%s",
        tableName, tableProperties());

    List<SimpleRecord> records =
        Lists.newArrayList(
            new SimpleRecord(1, "a"), new SimpleRecord(2, "b"), new SimpleRecord(3, "c"));
    Dataset<Row> df = spark.createDataFrame(records, SimpleRecord.class);
    df.coalesce(1).writeTo(tableName).append();

    assertEquals(
        "Should have 3 rows in 2 partitions",
        ImmutableList.of(row(1L, "a"), row(2L, "b"), row(3L, "c")),
        sql("SELECT * FROM %s ORDER BY id", tableName));

    sql("DELETE FROM %s WHERE id > 2", tableName);
    assertEquals(
        "Should have two rows in the second partition",
        ImmutableList.of(row(1L, "a"), row(2L, "b")),
        sql("SELECT * FROM %s ORDER BY id", tableName));

    sql("DELETE FROM %s WHERE id < 2", tableName);

    assertEquals(
        "Should have two rows in the second partition",
        ImmutableList.of(row(2L, "b")),
        sql("SELECT * FROM %s ORDER BY id", tableName));
  }

  @TestTemplate
  public void testDeleteFromWhereFalse() {
    sql(
        "CREATE TABLE %s (id bigint NOT NULL, data string) USING iceberg %s",
        tableName, tableProperties());
    sql("INSERT INTO TABLE %s VALUES (1, 'a'), (2, 'b'), (3, 'c')", tableName);

    assertEquals(
        "Should have expected rows",
        ImmutableList.of(row(1L, "a"), row(2L, "b"), row(3L, "c")),
        sql("SELECT * FROM %s ORDER BY id", tableName));

    Table table = validationCatalog.loadTable(tableIdent);
    assertThat(table.snapshots()).as("Should have 1 snapshot").hasSize(1);

    sql("DELETE FROM %s WHERE false", tableName);

    table.refresh();

    assertThat(table.snapshots()).as("Delete should not produce a new snapshot").hasSize(1);
  }

  @TestTemplate
  public void testTruncate() {
    sql(
        "CREATE TABLE %s (id bigint NOT NULL, data string) USING iceberg %s",
        tableName, tableProperties());
    sql("INSERT INTO TABLE %s VALUES (1, 'a'), (2, 'b'), (3, 'c')", tableName);

    assertEquals(
        "Should have expected rows",
        ImmutableList.of(row(1L, "a"), row(2L, "b"), row(3L, "c")),
        sql("SELECT * FROM %s ORDER BY id", tableName));

    Table table = validationCatalog.loadTable(tableIdent);
    assertThat(table.snapshots()).as("Should have 1 snapshot").hasSize(1);

    sql("TRUNCATE TABLE %s", tableName);

    assertEquals(
        "Should have expected rows",
        ImmutableList.of(),
        sql("SELECT * FROM %s ORDER BY id", tableName));
  }

  @TestTemplate
  public void testDeleteFromTablePartitionedByVarbinary() {
    sql(
        "CREATE TABLE %s (id bigint NOT NULL, data binary) USING iceberg PARTITIONED BY (data) %s",
        tableName, tableProperties());
    sql("INSERT INTO TABLE %s VALUES(1, X'e3bcd1'), (2, X'bcd1')", tableName);

    assertEquals(
        "Should have expected rows",
        ImmutableList.of(row(1L, new byte[] {-29, -68, -47}), row(2L, new byte[] {-68, -47})),
        sql("SELECT * FROM %s ORDER BY id", tableName));

    sql("DELETE FROM %s WHERE data = X'bcd1'", tableName);
    assertEquals(
        "Should have expected rows",
        ImmutableList.of(row(1L, new byte[] {-29, -68, -47})),
        sql("SELECT * FROM %s", tableName));
    assertEquals(
        "Should have expected rows",
        ImmutableList.of(row(1L, new byte[] {-29, -68, -47})),
        sql("SELECT * FROM %s where data = X'e3bcd1'", tableName));
  }

  @TestTemplate
  public void truncateWithDVs() throws NoSuchTableException {
    assumeThat(formatVersion).isGreaterThanOrEqualTo(3);

    sql(
        "CREATE TABLE %s (id bigint NOT NULL, data string) USING iceberg %s",
        tableName,
        tableProperties(
            ImmutableMap.of(
                TableProperties.DELETE_MODE, RowLevelOperationMode.MERGE_ON_READ.modeName())));
    List<SimpleRecord> records =
        ImmutableList.of(
            new SimpleRecord(1, "a"), new SimpleRecord(2, "b"), new SimpleRecord(3, "c"));
    Dataset<Row> df = spark.createDataFrame(records, SimpleRecord.class);
    df.coalesce(1).writeTo(tableName).append();

    assertThat(sql("SELECT * FROM %s ORDER BY id", tableName))
        .containsExactly(row(1L, "a"), row(2L, "b"), row(3L, "c"));

    sql("DELETE FROM %s WHERE id = 1", tableName);
    assertThat(validationCatalog.loadTable(tableIdent).currentSnapshot().summary())
        .containsEntry(SnapshotSummary.ADDED_DVS_PROP, "1")
        .containsEntry(SnapshotSummary.ADDED_POS_DELETES_PROP, "1");

    sql("DELETE FROM %s WHERE id = 2", tableName);
    // DVs have been merged into single file
    assertThat(validationCatalog.loadTable(tableIdent).currentSnapshot().summary())
        .containsEntry(SnapshotSummary.ADDED_DVS_PROP, "1")
        .containsEntry(SnapshotSummary.REMOVED_DVS_PROP, "1")
        .containsEntry(SnapshotSummary.ADDED_POS_DELETES_PROP, "2");

    assertThat(sql("SELECT * FROM %s ORDER BY id", tableName)).containsExactly(row(3L, "c"));

    sql("TRUNCATE TABLE %s", tableName);
    assertThat(validationCatalog.loadTable(tableIdent).currentSnapshot().summary())
        .containsEntry(SnapshotSummary.REMOVED_DVS_PROP, "1")
        .containsEntry(SnapshotSummary.REMOVED_DELETE_FILES_PROP, "1")
        .containsEntry(SnapshotSummary.REMOVED_POS_DELETES_PROP, "2");

    assertThat(sql("SELECT * FROM %s ORDER BY id", tableName)).isEmpty();
  }
}
