/*
 * 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.flink.maintenance.api;

import static org.apache.iceberg.flink.SimpleDataUtil.createRecord;
import static org.apache.iceberg.flink.maintenance.api.RewriteDataFiles.COMMIT_TASK_NAME;
import static org.apache.iceberg.flink.maintenance.api.RewriteDataFiles.PLANNER_TASK_NAME;
import static org.apache.iceberg.flink.maintenance.api.RewriteDataFiles.REWRITE_TASK_NAME;
import static org.apache.iceberg.flink.maintenance.operator.TableMaintenanceMetrics.ADDED_DATA_FILE_NUM_METRIC;
import static org.apache.iceberg.flink.maintenance.operator.TableMaintenanceMetrics.ADDED_DATA_FILE_SIZE_METRIC;
import static org.apache.iceberg.flink.maintenance.operator.TableMaintenanceMetrics.ERROR_COUNTER;
import static org.apache.iceberg.flink.maintenance.operator.TableMaintenanceMetrics.PLANNED_GROUPS_COUNTER;
import static org.apache.iceberg.flink.maintenance.operator.TableMaintenanceMetrics.REMOVED_DATA_FILE_NUM_METRIC;
import static org.apache.iceberg.flink.maintenance.operator.TableMaintenanceMetrics.REMOVED_DATA_FILE_SIZE_METRIC;
import static org.assertj.core.api.Assertions.assertThat;

import java.time.Instant;
import java.util.List;
import java.util.stream.StreamSupport;
import org.apache.flink.streaming.api.graph.StreamGraphGenerator;
import org.apache.iceberg.FileFormat;
import org.apache.iceberg.ManifestFiles;
import org.apache.iceberg.MetadataColumns;
import org.apache.iceberg.Schema;
import org.apache.iceberg.SnapshotRef;
import org.apache.iceberg.Table;
import org.apache.iceberg.expressions.Expressions;
import org.apache.iceberg.flink.SimpleDataUtil;
import org.apache.iceberg.flink.maintenance.operator.MetricsReporterFactoryForTests;
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.FieldSource;

class TestRewriteDataFiles extends MaintenanceTaskTestBase {

  private static final FileFormat[] FILE_FORMATS =
      new FileFormat[] {FileFormat.AVRO, FileFormat.PARQUET, FileFormat.ORC};

  @Test
  void testRewriteUnpartitioned() throws Exception {
    Table table = createTable();
    insert(table, 1, "a");
    insert(table, 2, "b");
    insert(table, 3, "c");
    insert(table, 4, "d");

    assertFileNum(table, 4, 0);

    appendRewriteDataFiles(
        RewriteDataFiles.builder()
            .parallelism(2)
            .deleteFileThreshold(10)
            .targetFileSizeBytes(1_000_000L)
            .maxFileGroupSizeBytes(10_000_000L)
            .maxFileSizeBytes(2_000_000L)
            .minFileSizeBytes(500_000L)
            .minInputFiles(2)
            .partialProgressEnabled(true)
            .partialProgressMaxCommits(1)
            .maxRewriteBytes(100_000L)
            .rewriteAll(false));

    runAndWaitForSuccess(infra.env(), infra.source(), infra.sink());

    assertFileNum(table, 1, 0);

    SimpleDataUtil.assertTableRecords(
        table,
        ImmutableList.of(
            createRecord(1, "a"),
            createRecord(2, "b"),
            createRecord(3, "c"),
            createRecord(4, "d")));
  }

  @ParameterizedTest
  @FieldSource("FILE_FORMATS")
  void testRewriteUnpartitionedPreserveLineage(FileFormat fileFormat) throws Exception {
    Table table = createTable(3, fileFormat);
    insert(table, 1, "a", fileFormat);
    insert(table, 2, "b", fileFormat);
    insert(table, 3, "c", fileFormat);
    insert(table, 4, "d", fileFormat);

    assertFileNum(table, 4, 0);

    appendRewriteDataFiles(
        RewriteDataFiles.builder()
            .parallelism(2)
            .deleteFileThreshold(10)
            .targetFileSizeBytes(1_000_000L)
            .maxFileGroupSizeBytes(10_000_000L)
            .maxFileSizeBytes(2_000_000L)
            .minFileSizeBytes(500_000L)
            .minInputFiles(2)
            .partialProgressEnabled(true)
            .partialProgressMaxCommits(1)
            .maxRewriteBytes(100_000L)
            .rewriteAll(false));

    runAndWaitForSuccess(infra.env(), infra.source(), infra.sink());

    assertFileNum(table, 1, 0);

    Schema schema = MetadataColumns.schemaWithRowLineage(table.schema());
    SimpleDataUtil.assertTableRecords(
        table,
        ImmutableList.of(
            SimpleDataUtil.createRecordWithRowId(1, "a", 0L, 1L),
            SimpleDataUtil.createRecordWithRowId(2, "b", 1L, 2L),
            SimpleDataUtil.createRecordWithRowId(3, "c", 2L, 3L),
            SimpleDataUtil.createRecordWithRowId(4, "d", 3L, 4L)),
        SnapshotRef.MAIN_BRANCH,
        schema);
  }

  @ParameterizedTest
  @FieldSource("FILE_FORMATS")
  void testRewriteTheSameFilePreserveLineage(FileFormat fileFormat) throws Exception {
    Table table = createTable(3, fileFormat);
    insert(table, 1, "a", fileFormat);
    insert(table, 2, "b", fileFormat);
    // Create a file with two lines of data to verify that the rowid is read correctly.
    insert(
        table,
        ImmutableList.of(SimpleDataUtil.createRecord(3, "c"), SimpleDataUtil.createRecord(4, "d")),
        fileFormat);

    assertFileNum(table, 3, 0);

    appendRewriteDataFiles(
        RewriteDataFiles.builder()
            .parallelism(2)
            .deleteFileThreshold(10)
            .targetFileSizeBytes(1_000_000L)
            .maxFileGroupSizeBytes(10_000_000L)
            .maxFileSizeBytes(2_000_000L)
            .minFileSizeBytes(500_000L)
            .minInputFiles(2)
            .partialProgressEnabled(true)
            .partialProgressMaxCommits(1)
            .maxRewriteBytes(100_000L)
            .rewriteAll(false));

    runAndWaitForSuccess(infra.env(), infra.source(), infra.sink());

    assertFileNum(table, 1, 0);

    Schema schema = MetadataColumns.schemaWithRowLineage(table.schema());
    SimpleDataUtil.assertTableRecords(
        table,
        ImmutableList.of(
            SimpleDataUtil.createRecordWithRowId(1, "a", 0L, 1L),
            SimpleDataUtil.createRecordWithRowId(2, "b", 1L, 2L),
            // The Ids 3 and 4 come from the same file, so the last updated sequence number should
            // be the same.
            SimpleDataUtil.createRecordWithRowId(3, "c", 2L, 3L),
            SimpleDataUtil.createRecordWithRowId(4, "d", 3L, 3L)),
        SnapshotRef.MAIN_BRANCH,
        schema);
  }

  @ParameterizedTest
  @FieldSource("FILE_FORMATS")
  void testRewritePartitionedPreserveLineage(FileFormat fileFormat) throws Exception {
    Table table = createPartitionedTable(3, fileFormat);
    insertPartitioned(table, 1, "p1", fileFormat);
    insertPartitioned(table, 2, "p1", fileFormat);
    insertPartitioned(table, 3, "p2", fileFormat);
    insertPartitioned(table, 4, "p2", fileFormat);

    assertFileNum(table, 4, 0);

    appendRewriteDataFiles();

    runAndWaitForSuccess(infra.env(), infra.source(), infra.sink());

    assertFileNum(table, 2, 0);

    Schema schema = MetadataColumns.schemaWithRowLineage(table.schema());
    SimpleDataUtil.assertTableRecords(
        table,
        ImmutableList.of(
            SimpleDataUtil.createRecordWithRowId(1, "p1", 0L, 1L),
            SimpleDataUtil.createRecordWithRowId(2, "p1", 1L, 2L),
            SimpleDataUtil.createRecordWithRowId(3, "p2", 2L, 3L),
            SimpleDataUtil.createRecordWithRowId(4, "p2", 3L, 4L)),
        SnapshotRef.MAIN_BRANCH,
        schema);
  }

  @Test
  void testRewritePartitioned() throws Exception {
    Table table = createPartitionedTable();
    insertPartitioned(table, 1, "p1");
    insertPartitioned(table, 2, "p1");
    insertPartitioned(table, 3, "p2");
    insertPartitioned(table, 4, "p2");

    assertFileNum(table, 4, 0);

    appendRewriteDataFiles();

    runAndWaitForSuccess(infra.env(), infra.source(), infra.sink());

    assertFileNum(table, 2, 0);

    SimpleDataUtil.assertTableRecords(
        table,
        ImmutableList.of(
            createRecord(1, "p1"),
            createRecord(2, "p1"),
            createRecord(3, "p2"),
            createRecord(4, "p2")));
  }

  @Test
  void testPlannerFailure() throws Exception {
    Table table = createTable();
    insert(table, 1, "a");
    insert(table, 2, "b");

    assertFileNum(table, 2, 0);

    appendRewriteDataFiles();

    runAndWaitForFailure(infra.env(), infra.source(), infra.sink());

    // Check the metrics. The first task should be successful, but the second one should fail. This
    // should be represented in the counters.
    MetricsReporterFactoryForTests.assertCounters(
        new ImmutableMap.Builder<List<String>, Long>()
            .put(
                ImmutableList.of(
                    PLANNER_TASK_NAME + "[0]",
                    DUMMY_TABLE_NAME,
                    DUMMY_TASK_NAME,
                    "0",
                    ERROR_COUNTER),
                1L)
            .put(
                ImmutableList.of(
                    PLANNER_TASK_NAME + "[0]",
                    DUMMY_TABLE_NAME,
                    DUMMY_TASK_NAME,
                    "0",
                    PLANNED_GROUPS_COUNTER),
                1L)
            .put(
                ImmutableList.of(
                    REWRITE_TASK_NAME + "[0]",
                    DUMMY_TABLE_NAME,
                    DUMMY_TASK_NAME,
                    "0",
                    ERROR_COUNTER),
                0L)
            .put(
                ImmutableList.of(
                    COMMIT_TASK_NAME + "[0]",
                    DUMMY_TABLE_NAME,
                    DUMMY_TASK_NAME,
                    "0",
                    ERROR_COUNTER),
                0L)
            .put(
                ImmutableList.of(
                    COMMIT_TASK_NAME + "[0]",
                    DUMMY_TABLE_NAME,
                    DUMMY_TASK_NAME,
                    "0",
                    ADDED_DATA_FILE_NUM_METRIC),
                1L)
            .put(
                ImmutableList.of(
                    COMMIT_TASK_NAME + "[0]",
                    DUMMY_TABLE_NAME,
                    DUMMY_TASK_NAME,
                    "0",
                    ADDED_DATA_FILE_SIZE_METRIC),
                -1L)
            .put(
                ImmutableList.of(
                    COMMIT_TASK_NAME + "[0]",
                    DUMMY_TABLE_NAME,
                    DUMMY_TASK_NAME,
                    "0",
                    REMOVED_DATA_FILE_NUM_METRIC),
                2L)
            .put(
                ImmutableList.of(
                    COMMIT_TASK_NAME + "[0]",
                    DUMMY_TABLE_NAME,
                    DUMMY_TASK_NAME,
                    "0",
                    REMOVED_DATA_FILE_SIZE_METRIC),
                -1L)
            .build());
  }

  @Test
  void testUidAndSlotSharingGroup() {
    createTable();

    RewriteDataFiles.builder()
        .slotSharingGroup(SLOT_SHARING_GROUP)
        .uidSuffix(UID_SUFFIX)
        .append(
            infra.triggerStream(),
            DUMMY_TABLE_NAME,
            DUMMY_TASK_NAME,
            0,
            tableLoader(),
            "OTHER",
            "OTHER",
            1)
        .sinkTo(infra.sink());

    checkUidsAreSet(infra.env(), UID_SUFFIX);
    checkSlotSharingGroupsAreSet(infra.env(), SLOT_SHARING_GROUP);
  }

  @Test
  void testUidAndSlotSharingGroupUnset() {
    createTable();

    RewriteDataFiles.builder()
        .append(
            infra.triggerStream(),
            DUMMY_TABLE_NAME,
            DUMMY_TASK_NAME,
            0,
            tableLoader(),
            UID_SUFFIX,
            null,
            1)
        .sinkTo(infra.sink());

    checkUidsAreSet(infra.env(), null);
    checkSlotSharingGroupsAreSet(infra.env(), null);
  }

  @Test
  void testMetrics() throws Exception {
    Table table = createTable();
    insert(table, 1, "a");
    insert(table, 2, "b");

    assertFileNum(table, 2, 0);

    appendRewriteDataFiles();

    runAndWaitForSuccess(infra.env(), infra.source(), infra.sink());

    // Check the metrics
    MetricsReporterFactoryForTests.assertCounters(
        new ImmutableMap.Builder<List<String>, Long>()
            .put(
                ImmutableList.of(
                    PLANNER_TASK_NAME + "[0]",
                    DUMMY_TABLE_NAME,
                    DUMMY_TASK_NAME,
                    "0",
                    ERROR_COUNTER),
                0L)
            .put(
                ImmutableList.of(
                    PLANNER_TASK_NAME + "[0]",
                    DUMMY_TABLE_NAME,
                    DUMMY_TASK_NAME,
                    "0",
                    PLANNED_GROUPS_COUNTER),
                1L)
            .put(
                ImmutableList.of(
                    REWRITE_TASK_NAME + "[0]",
                    DUMMY_TABLE_NAME,
                    DUMMY_TASK_NAME,
                    "0",
                    ERROR_COUNTER),
                0L)
            .put(
                ImmutableList.of(
                    COMMIT_TASK_NAME + "[0]",
                    DUMMY_TABLE_NAME,
                    DUMMY_TASK_NAME,
                    "0",
                    ERROR_COUNTER),
                0L)
            .put(
                ImmutableList.of(
                    COMMIT_TASK_NAME + "[0]",
                    DUMMY_TABLE_NAME,
                    DUMMY_TASK_NAME,
                    "0",
                    ADDED_DATA_FILE_NUM_METRIC),
                1L)
            .put(
                ImmutableList.of(
                    COMMIT_TASK_NAME + "[0]",
                    DUMMY_TABLE_NAME,
                    DUMMY_TASK_NAME,
                    "0",
                    ADDED_DATA_FILE_SIZE_METRIC),
                -1L)
            .put(
                ImmutableList.of(
                    COMMIT_TASK_NAME + "[0]",
                    DUMMY_TABLE_NAME,
                    DUMMY_TASK_NAME,
                    "0",
                    REMOVED_DATA_FILE_NUM_METRIC),
                2L)
            .put(
                ImmutableList.of(
                    COMMIT_TASK_NAME + "[0]",
                    DUMMY_TABLE_NAME,
                    DUMMY_TASK_NAME,
                    "0",
                    REMOVED_DATA_FILE_SIZE_METRIC),
                -1L)
            .build());
  }

  @Test
  void testV2Table() throws Exception {
    Table table = createTableWithDelete();
    update(table, 1, null, "a", "b");
    update(table, 1, "b", "c");

    assertFileNum(table, 2, 3);
    SimpleDataUtil.assertTableRecords(table, ImmutableList.of(createRecord(1, "c")));

    appendRewriteDataFiles();

    runAndWaitForSuccess(infra.env(), infra.source(), infra.sink());

    assertFileNum(table, 1, 1);

    SimpleDataUtil.assertTableRecords(table, ImmutableList.of(createRecord(1, "c")));

    // Check the metrics
    MetricsReporterFactoryForTests.assertCounters(
        new ImmutableMap.Builder<List<String>, Long>()
            .put(
                ImmutableList.of(
                    PLANNER_TASK_NAME + "[0]",
                    DUMMY_TABLE_NAME,
                    DUMMY_TASK_NAME,
                    "0",
                    ERROR_COUNTER),
                0L)
            .put(
                ImmutableList.of(
                    PLANNER_TASK_NAME + "[0]",
                    DUMMY_TABLE_NAME,
                    DUMMY_TASK_NAME,
                    "0",
                    PLANNED_GROUPS_COUNTER),
                1L)
            .put(
                ImmutableList.of(
                    REWRITE_TASK_NAME + "[0]",
                    DUMMY_TABLE_NAME,
                    DUMMY_TASK_NAME,
                    "0",
                    ERROR_COUNTER),
                0L)
            .put(
                ImmutableList.of(
                    COMMIT_TASK_NAME + "[0]",
                    DUMMY_TABLE_NAME,
                    DUMMY_TASK_NAME,
                    "0",
                    ERROR_COUNTER),
                0L)
            .put(
                ImmutableList.of(
                    COMMIT_TASK_NAME + "[0]",
                    DUMMY_TABLE_NAME,
                    DUMMY_TASK_NAME,
                    "0",
                    ADDED_DATA_FILE_NUM_METRIC),
                1L)
            .put(
                ImmutableList.of(
                    COMMIT_TASK_NAME + "[0]",
                    DUMMY_TABLE_NAME,
                    DUMMY_TASK_NAME,
                    "0",
                    ADDED_DATA_FILE_SIZE_METRIC),
                -1L)
            .put(
                ImmutableList.of(
                    COMMIT_TASK_NAME + "[0]",
                    DUMMY_TABLE_NAME,
                    DUMMY_TASK_NAME,
                    "0",
                    REMOVED_DATA_FILE_NUM_METRIC),
                2L)
            .put(
                ImmutableList.of(
                    COMMIT_TASK_NAME + "[0]",
                    DUMMY_TABLE_NAME,
                    DUMMY_TASK_NAME,
                    "0",
                    REMOVED_DATA_FILE_SIZE_METRIC),
                -1L)
            .build());
  }

  @Test
  void testRewriteWithFilter() throws Exception {
    Table table = createTable();
    insert(table, 1, "a");
    insert(table, 2, "b");
    insert(table, 3, "c");
    insert(table, 4, "d");

    assertFileNum(table, 4, 0);

    appendRewriteDataFiles(
        RewriteDataFiles.builder()
            .parallelism(2)
            .deleteFileThreshold(10)
            .targetFileSizeBytes(1_000_000L)
            .maxFileGroupSizeBytes(10_000_000L)
            .maxFileSizeBytes(2_000_000L)
            .minFileSizeBytes(500_000L)
            .minInputFiles(2)
            // Only rewrite data files where id is 1 or 2 for testing rewrite
            .filter(Expressions.in("id", 1, 2))
            .partialProgressEnabled(true)
            .partialProgressMaxCommits(1)
            .maxRewriteBytes(100_000L)
            .rewriteAll(false));

    runAndWaitForSuccess(infra.env(), infra.source(), infra.sink());

    // There is four files, only id is 1 and 2 will be rewritten. so expect 3 files.
    assertFileNum(table, 3, 0);

    SimpleDataUtil.assertTableRecords(
        table,
        ImmutableList.of(
            createRecord(1, "a"),
            createRecord(2, "b"),
            createRecord(3, "c"),
            createRecord(4, "d")));
  }

  /**
   * By verifying that the creation time of the data content in the builder is later than the
   * creation time of the filter condition — if the filter condition is actually created in the
   * planner, then all files can be compacted; otherwise, not all files can be compacted — we can
   * confirm whether the filter condition is actually created in the planner.
   */
  @Test
  void testRewriteWithFilterSupplier() throws Exception {
    Table table = createTable();

    appendRewriteDataFiles(
        RewriteDataFiles.builder()
            .parallelism(2)
            .deleteFileThreshold(10)
            .targetFileSizeBytes(1_000_000L)
            .maxFileGroupSizeBytes(10_000_000L)
            .maxFileSizeBytes(2_000_000L)
            .minFileSizeBytes(500_000L)
            .minInputFiles(2)
            // Rewrite data files where id is less than current timestamp in planner
            .filter(() -> Expressions.lessThan("id", (int) Instant.now().getEpochSecond()))
            .partialProgressEnabled(true)
            .partialProgressMaxCommits(1)
            .maxRewriteBytes(100_000L)
            .rewriteAll(false));

    insert(table, 1, "a");
    insert(table, 2, "b");
    insert(table, 3, "c");

    int epochSecond = (int) Instant.now().getEpochSecond();
    insert(table, epochSecond, "d");

    assertFileNum(table, 4, 0);

    Thread.sleep(1_000L);
    runAndWaitForSuccess(infra.env(), infra.source(), infra.sink());

    // There is four files, only id is less than current timestamp will be rewritten. so expect 2
    // files.
    assertFileNum(table, 1, 0);

    SimpleDataUtil.assertTableRecords(
        table,
        ImmutableList.of(
            createRecord(1, "a"),
            createRecord(2, "b"),
            createRecord(3, "c"),
            createRecord(epochSecond, "d")));
  }

  @Test
  void testBranch() throws Exception {
    Table table = createTable();
    insert(table, 1, "a");
    insert(table, 2, "b");

    // Create branch based on above inserts
    String branchName = "test-branch";
    table.manageSnapshots().createBranch(branchName).commit();

    // Insert another file on main only (main has 3 files, branch stays at 2)
    insert(table, 3, "c");

    appendRewriteDataFiles(RewriteDataFiles.builder().rewriteAll(true).branch(branchName));

    runAndWaitForSuccess(infra.env(), infra.source(), infra.sink());

    table.refresh();

    // Branch should be compacted from 2 files to 1
    assertThat(
            table.snapshot(branchName).dataManifests(table.io()).stream()
                .flatMap(
                    m ->
                        StreamSupport.stream(
                            ManifestFiles.read(m, table.io(), table.specs()).spliterator(), false))
                .count())
        .isEqualTo(1);
    SimpleDataUtil.assertTableRecords(
        table, ImmutableList.of(createRecord(1, "a"), createRecord(2, "b")), branchName);

    // Main should be untouched with 3 files
    assertFileNum(table, 3, 0);
  }

  private void appendRewriteDataFiles() {
    appendRewriteDataFiles(RewriteDataFiles.builder().rewriteAll(true));
  }

  private void appendRewriteDataFiles(RewriteDataFiles.Builder builder) {
    builder
        .append(
            infra.triggerStream(),
            DUMMY_TABLE_NAME,
            DUMMY_TASK_NAME,
            0,
            tableLoader(),
            UID_SUFFIX,
            StreamGraphGenerator.DEFAULT_SLOT_SHARING_GROUP,
            1)
        .sinkTo(infra.sink());
  }

  private static void assertFileNum(
      Table table, int expectedDataFileNum, int expectedDeleteFileNum) {
    table.refresh();
    assertThat(
            table.currentSnapshot().dataManifests(table.io()).stream()
                .flatMap(
                    m ->
                        StreamSupport.stream(
                            ManifestFiles.read(m, table.io(), table.specs()).spliterator(), false))
                .count())
        .isEqualTo(expectedDataFileNum);
    assertThat(
            table.currentSnapshot().deleteManifests(table.io()).stream()
                .flatMap(
                    m ->
                        StreamSupport.stream(
                            ManifestFiles.readDeleteManifest(m, table.io(), table.specs())
                                .spliterator(),
                            false))
                .count())
        .isEqualTo(expectedDeleteFileNum);
  }
}
