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

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

import java.io.IOException;
import java.nio.file.Path;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import org.apache.flink.core.execution.JobClient;
import org.apache.flink.runtime.jobgraph.JobVertex;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.graph.StreamGraph;
import org.apache.flink.table.data.RowData;
import org.apache.flink.types.Row;
import org.apache.iceberg.DataFile;
import org.apache.iceberg.DeleteFile;
import org.apache.iceberg.FileFormat;
import org.apache.iceberg.ManifestFile;
import org.apache.iceberg.ManifestFiles;
import org.apache.iceberg.ManifestReader;
import org.apache.iceberg.PartitionSpec;
import org.apache.iceberg.Snapshot;
import org.apache.iceberg.Table;
import org.apache.iceberg.TableProperties;
import org.apache.iceberg.actions.SizeBasedFileRewritePlanner;
import org.apache.iceberg.data.GenericAppenderHelper;
import org.apache.iceberg.flink.FlinkWriteOptions;
import org.apache.iceberg.flink.MiniFlinkClusterExtension;
import org.apache.iceberg.flink.SimpleDataUtil;
import org.apache.iceberg.flink.TestFixtures;
import org.apache.iceberg.flink.maintenance.api.DeleteOrphanFilesConfig;
import org.apache.iceberg.flink.maintenance.api.ExpireSnapshotsConfig;
import org.apache.iceberg.flink.maintenance.api.LockConfig;
import org.apache.iceberg.flink.maintenance.api.RewriteDataFilesConfig;
import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
import org.apache.iceberg.relocated.com.google.common.collect.Lists;
import org.apache.iceberg.relocated.com.google.common.collect.Maps;
import org.apache.iceberg.util.ContentFileUtil;
import org.awaitility.Awaitility;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.io.TempDir;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.FieldSource;

class TestIcebergSinkTableMaintenance extends TestFlinkIcebergSinkBase {
  private static final String[] LOCK_TYPES = new String[] {LockConfig.JdbcLockConfig.JDBC, ""};

  @TempDir private Path tempDir;

  private Map<String, String> flinkConf;

  @BeforeEach
  void before() throws IOException {
    this.flinkConf = Maps.newHashMap();

    table =
        CATALOG_EXTENSION
            .catalog()
            .createTable(
                TestFixtures.TABLE_IDENTIFIER,
                SimpleDataUtil.SCHEMA,
                PartitionSpec.unpartitioned(),
                Maps.newHashMap());

    env =
        StreamExecutionEnvironment.getExecutionEnvironment(
                MiniFlinkClusterExtension.DISABLE_CLASSLOADER_CHECK_CONFIG)
            .enableCheckpointing(100);

    tableLoader = CATALOG_EXTENSION.tableLoader();
  }

  @ParameterizedTest(name = "lockType = {0}")
  @FieldSource("LOCK_TYPES")
  public void testCompactFileE2e(String lockType) throws Exception {
    setupLockConfig(lockType);
    setupCompactionConfig();
    List<Row> rows = Lists.newArrayList(Row.of(1, "hello"), Row.of(2, "world"), Row.of(3, "foo"));
    DataStream<RowData> dataStream =
        env.addSource(createBoundedSource(rows), ROW_TYPE_INFO)
            .map(CONVERTER::toInternal, FlinkCompatibilityUtil.toTypeInfo(SimpleDataUtil.ROW_TYPE));

    IcebergSink.forRowData(dataStream)
        .setAll(flinkConf)
        .table(table)
        .tableLoader(tableLoader)
        .append();

    env.execute("Test Iceberg Compaction DataStream");

    table.refresh();
    // check the data file count after compact
    List<DataFile> afterCompactDataFiles = getDataFiles(table.currentSnapshot(), table);
    assertThat(afterCompactDataFiles).hasSize(1);

    // check the data file count before compact
    List<DataFile> preCompactDataFiles =
        getDataFiles(table.snapshot(table.currentSnapshot().parentId()), table);
    assertThat(preCompactDataFiles).hasSize(3);
  }

  private List<DataFile> getDataFiles(Snapshot snapshot, Table table) throws IOException {
    List<DataFile> dataFiles = Lists.newArrayList();
    for (ManifestFile dataManifest : snapshot.dataManifests(table.io())) {
      try (ManifestReader<DataFile> reader =
          ManifestFiles.read(dataManifest, table.io(), table.specs())) {
        reader.iterator().forEachRemaining(dataFiles::add);
      }
    }

    return dataFiles;
  }

  @ParameterizedTest(name = "lockType = {0}")
  @FieldSource("LOCK_TYPES")
  public void testTableMaintenanceOperatorAdded(String lockType) {
    setupLockConfig(lockType);
    setupCompactionConfig();
    List<Row> rows = Lists.newArrayList(Row.of(1, "hello"), Row.of(2, "world"), Row.of(3, "foo"));
    DataStream<RowData> dataStream =
        env.addSource(createBoundedSource(rows), ROW_TYPE_INFO)
            .map(CONVERTER::toInternal, FlinkCompatibilityUtil.toTypeInfo(SimpleDataUtil.ROW_TYPE));

    IcebergSink.forRowData(dataStream)
        .table(table)
        .tableLoader(tableLoader)
        .setAll(flinkConf)
        .append();

    boolean containRewrite = false;
    StreamGraph streamGraph = env.getStreamGraph();
    for (JobVertex vertex : streamGraph.getJobGraph().getVertices()) {
      if (vertex.getName().contains("Rewrite")) {
        containRewrite = true;
        break;
      }
    }

    assertThat(containRewrite).isTrue();
  }

  private void setupCompactionConfig() {
    flinkConf.put(FlinkWriteOptions.COMPACTION_ENABLE.key(), "true");
    flinkConf.put(RewriteDataFilesConfig.SCHEDULE_ON_DATA_FILE_SIZE, "1");
    flinkConf.put(RewriteDataFilesConfig.PREFIX + SizeBasedFileRewritePlanner.REWRITE_ALL, "true");
  }

  private void setupExpireSnapshotsConfig() {
    flinkConf.put(FlinkWriteOptions.EXPIRE_SNAPSHOTS_ENABLE.key(), "true");
    flinkConf.put(ExpireSnapshotsConfig.RETAIN_LAST, "5");
  }

  private void setupDeleteOrphanFilesConfig() {
    flinkConf.put(FlinkWriteOptions.DELETE_ORPHAN_FILES_ENABLE.key(), "true");
    flinkConf.put(DeleteOrphanFilesConfig.MIN_AGE_SECONDS, "86400");
  }

  private void setupConvertEqualityDeletesConfig() {
    flinkConf.put(FlinkWriteOptions.CONVERT_EQUALITY_DELETES_ENABLE.key(), "true");
  }

  private void upgradeToFormatV3() {
    table.updateProperties().set(TableProperties.FORMAT_VERSION, "3").commit();
    table.refresh();
  }

  private void setupLockConfig(String lockType) {
    if (lockType.equals(LockConfig.JdbcLockConfig.JDBC)) {
      flinkConf.put(LockConfig.LOCK_TYPE_OPTION.key(), LockConfig.JdbcLockConfig.JDBC);
      flinkConf.put(
          LockConfig.JdbcLockConfig.JDBC_URI_OPTION.key(),
          "jdbc:sqlite:file::memory:?ic" + UUID.randomUUID().toString().replace("-", ""));
      flinkConf.put(LockConfig.LOCK_ID_OPTION.key(), "test-lock-id");
      flinkConf.put(LockConfig.JdbcLockConfig.JDBC_INIT_LOCK_TABLE_OPTION.key(), "true");
    } else {
      flinkConf.put(LockConfig.LOCK_TYPE_OPTION.key(), "");
    }
  }

  @ParameterizedTest(name = "lockType = {0}")
  @FieldSource("LOCK_TYPES")
  public void testExpireSnapshotsEnabled(String lockType) {
    setupLockConfig(lockType);
    setupExpireSnapshotsConfig();

    List<Row> rows = Lists.newArrayList(Row.of(1, "hello"), Row.of(2, "world"), Row.of(3, "foo"));
    DataStream<RowData> dataStream =
        env.addSource(createBoundedSource(rows), ROW_TYPE_INFO)
            .map(CONVERTER::toInternal, FlinkCompatibilityUtil.toTypeInfo(SimpleDataUtil.ROW_TYPE));

    IcebergSink.forRowData(dataStream)
        .table(table)
        .tableLoader(tableLoader)
        .setAll(flinkConf)
        .append();

    StreamGraph streamGraph = env.getStreamGraph();
    boolean containExpire = false;
    for (JobVertex vertex : streamGraph.getJobGraph().getVertices()) {
      if (vertex.getName().contains("Expire")) {
        containExpire = true;
        break;
      }
    }

    assertThat(containExpire).isTrue();
  }

  @ParameterizedTest(name = "lockType = {0}")
  @FieldSource("LOCK_TYPES")
  public void testDeleteOrphanFilesEnabled(String lockType) {
    setupLockConfig(lockType);
    setupDeleteOrphanFilesConfig();

    List<Row> rows = Lists.newArrayList(Row.of(1, "hello"), Row.of(2, "world"), Row.of(3, "foo"));
    DataStream<RowData> dataStream =
        env.addSource(createBoundedSource(rows), ROW_TYPE_INFO)
            .map(CONVERTER::toInternal, FlinkCompatibilityUtil.toTypeInfo(SimpleDataUtil.ROW_TYPE));

    IcebergSink.forRowData(dataStream)
        .table(table)
        .tableLoader(tableLoader)
        .setAll(flinkConf)
        .append();

    StreamGraph streamGraph = env.getStreamGraph();
    boolean containOrphan = false;
    for (JobVertex vertex : streamGraph.getJobGraph().getVertices()) {
      if (vertex.getName().contains("Orphan")) {
        containOrphan = true;
        break;
      }
    }

    assertThat(containOrphan).isTrue();
  }

  @ParameterizedTest(name = "lockType = {0}")
  @FieldSource("LOCK_TYPES")
  public void testAllMaintenanceTasksCombined(String lockType) {
    setupLockConfig(lockType);
    setupCompactionConfig();
    setupExpireSnapshotsConfig();
    setupDeleteOrphanFilesConfig();

    List<Row> rows = Lists.newArrayList(Row.of(1, "hello"), Row.of(2, "world"), Row.of(3, "foo"));
    DataStream<RowData> dataStream =
        env.addSource(createBoundedSource(rows), ROW_TYPE_INFO)
            .map(CONVERTER::toInternal, FlinkCompatibilityUtil.toTypeInfo(SimpleDataUtil.ROW_TYPE));

    IcebergSink.forRowData(dataStream)
        .table(table)
        .tableLoader(tableLoader)
        .setAll(flinkConf)
        .append();

    StreamGraph streamGraph = env.getStreamGraph();
    boolean containRewrite = false;
    boolean containExpire = false;
    boolean containOrphan = false;
    for (JobVertex vertex : streamGraph.getJobGraph().getVertices()) {
      if (vertex.getName().contains("Rewrite")) {
        containRewrite = true;
      }

      if (vertex.getName().contains("Expire")) {
        containExpire = true;
      }

      if (vertex.getName().contains("Orphan")) {
        containOrphan = true;
      }
    }

    assertThat(containRewrite).isTrue();
    assertThat(containExpire).isTrue();
    assertThat(containOrphan).isTrue();
  }

  @ParameterizedTest(name = "lockType = {0}")
  @FieldSource("LOCK_TYPES")
  public void testAllMaintenanceE2e(String lockType) throws Exception {
    setupLockConfig(lockType);

    Map<String, String> compactionConfig = Maps.newHashMap();
    compactionConfig.put(RewriteDataFilesConfig.SCHEDULE_ON_DATA_FILE_SIZE, "1");
    compactionConfig.put(
        RewriteDataFilesConfig.PREFIX + SizeBasedFileRewritePlanner.REWRITE_ALL, "true");

    Map<String, String> expireConfig = Maps.newHashMap();
    expireConfig.put(ExpireSnapshotsConfig.RETAIN_LAST, "5");

    Map<String, String> orphanConfig = Maps.newHashMap();
    orphanConfig.put(DeleteOrphanFilesConfig.MIN_AGE_SECONDS, "86400");

    List<Row> rows = Lists.newArrayList(Row.of(1, "hello"), Row.of(2, "world"), Row.of(3, "foo"));
    DataStream<RowData> dataStream =
        env.addSource(createBoundedSource(rows), ROW_TYPE_INFO)
            .map(CONVERTER::toInternal, FlinkCompatibilityUtil.toTypeInfo(SimpleDataUtil.ROW_TYPE));

    IcebergSink.forRowData(dataStream)
        .setAll(flinkConf)
        .table(table)
        .tableLoader(tableLoader)
        .rewriteDataFiles(compactionConfig)
        .expireSnapshots(expireConfig)
        .deleteOrphanFiles(orphanConfig)
        .append();

    env.execute("Test All Maintenance E2E");

    table.refresh();
    // Compaction should have merged the 3 data files into 1
    List<DataFile> afterCompactDataFiles = getDataFiles(table.currentSnapshot(), table);
    assertThat(afterCompactDataFiles).hasSize(1);

    List<DataFile> preCompactDataFiles =
        getDataFiles(table.snapshot(table.currentSnapshot().parentId()), table);
    assertThat(preCompactDataFiles).hasSize(3);
  }

  @ParameterizedTest(name = "lockType = {0}")
  @FieldSource("LOCK_TYPES")
  public void testConvertEqualityDeletesOperatorAdded(String lockType) {
    setupLockConfig(lockType);
    setupConvertEqualityDeletesConfig();
    upgradeToFormatV3();

    List<Row> rows = Lists.newArrayList(Row.of(1, "hello"));
    DataStream<RowData> dataStream =
        env.addSource(createBoundedSource(rows), ROW_TYPE_INFO)
            .map(CONVERTER::toInternal, FlinkCompatibilityUtil.toTypeInfo(SimpleDataUtil.ROW_TYPE));

    IcebergSink.forRowData(dataStream)
        .table(table)
        .tableLoader(tableLoader)
        .upsert(true)
        .equalityFieldColumns(Lists.newArrayList("id"))
        .setAll(flinkConf)
        .append();

    StreamGraph streamGraph = env.getStreamGraph();
    boolean containsConvert = false;
    for (JobVertex vertex : streamGraph.getJobGraph().getVertices()) {
      if (vertex.getName().contains("EqConvert")) {
        containsConvert = true;
        break;
      }
    }

    assertThat(containsConvert).isTrue();
  }

  @ParameterizedTest(name = "lockType = {0}")
  @FieldSource("LOCK_TYPES")
  public void testConvertEqualityDeletesRequiresEqualityFields(String lockType) {
    setupLockConfig(lockType);
    setupConvertEqualityDeletesConfig();

    List<Row> rows = Lists.newArrayList(Row.of(1, "hello"));
    DataStream<RowData> dataStream =
        env.addSource(createBoundedSource(rows), ROW_TYPE_INFO)
            .map(CONVERTER::toInternal, FlinkCompatibilityUtil.toTypeInfo(SimpleDataUtil.ROW_TYPE));

    // No equality field columns and no identifier fields on the schema, so there is nothing to
    // convert.
    assertThatThrownBy(
            () ->
                IcebergSink.forRowData(dataStream)
                    .table(table)
                    .tableLoader(tableLoader)
                    .setAll(flinkConf)
                    .append())
        .isInstanceOf(IllegalStateException.class)
        .hasMessageContaining("Equality field columns must be set");
  }

  @ParameterizedTest(name = "lockType = {0}")
  @FieldSource("LOCK_TYPES")
  public void testConvertEqualityDeletesE2e(String lockType) throws Exception {
    setupLockConfig(lockType);
    setupConvertEqualityDeletesConfig();
    upgradeToFormatV3();

    // Pre-existing row on main. The sink then upserts the same key, writing an equality delete that
    // removes this row, which the converter resolves to a deletion vector in a single cycle.
    new GenericAppenderHelper(table, FileFormat.PARQUET, tempDir)
        .appendToTable(ImmutableList.of(SimpleDataUtil.createRecord(1, "aaa")));

    List<Row> rows = Lists.newArrayList(Row.of(1, "bbb"));
    DataStream<RowData> dataStream =
        env.addSource(createBoundedSource(rows), ROW_TYPE_INFO)
            .map(CONVERTER::toInternal, FlinkCompatibilityUtil.toTypeInfo(SimpleDataUtil.ROW_TYPE));

    IcebergSink.forRowData(dataStream)
        .table(table)
        .tableLoader(tableLoader)
        .upsert(true)
        .equalityFieldColumns(Lists.newArrayList("id"))
        .setAll(flinkConf)
        .append();

    JobClient jobClient = env.executeAsync("Test Convert Equality Deletes E2E");
    try {
      Awaitility.await()
          .atMost(Duration.ofMinutes(2))
          .pollInterval(Duration.ofMillis(200))
          .untilAsserted(() -> assertThat(dvCountOnMain(table)).isEqualTo(1));
      SimpleDataUtil.assertTableRecords(
          table, ImmutableList.of(SimpleDataUtil.createRecord(1, "bbb")));
    } finally {
      jobClient.cancel();
    }
  }

  private static long dvCountOnMain(Table table) throws IOException {
    table.refresh();
    if (table.currentSnapshot() == null) {
      return 0;
    }

    long count = 0;
    for (ManifestFile manifest : table.currentSnapshot().deleteManifests(table.io())) {
      try (ManifestReader<DeleteFile> reader =
          ManifestFiles.readDeleteManifest(manifest, table.io(), table.specs())) {
        for (DeleteFile file : reader) {
          if (ContentFileUtil.isDV(file)) {
            count++;
          }
        }
      }
    }

    return count;
  }
}
