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

import org.apache.hudi.client.HoodieFlinkWriteClient;
import org.apache.hudi.client.common.HoodieFlinkEngineContext;
import org.apache.hudi.client.model.CommitTimeFlinkRecordMerger;
import org.apache.hudi.client.model.EventTimeFlinkRecordMerger;
import org.apache.hudi.client.model.PartialUpdateFlinkRecordMerger;
import org.apache.hudi.client.transaction.lock.FileSystemBasedLockProvider;
import org.apache.hudi.common.config.HoodieMetadataConfig;
import org.apache.hudi.common.config.RecordMergeMode;
import org.apache.hudi.common.model.DefaultHoodieRecordPayload;
import org.apache.hudi.common.model.EventTimeAvroPayload;
import org.apache.hudi.common.model.HoodieTableType;
import org.apache.hudi.common.model.OverwriteWithLatestAvroPayload;
import org.apache.hudi.common.model.PartialUpdateAvroPayload;
import org.apache.hudi.common.model.WriteConcurrencyMode;
import org.apache.hudi.common.table.HoodieTableConfig;
import org.apache.hudi.common.table.HoodieTableMetaClient;
import org.apache.hudi.common.table.HoodieTableVersion;
import org.apache.hudi.common.table.marker.MarkerType;
import org.apache.hudi.config.HoodieWriteConfig;
import org.apache.hudi.configuration.FlinkOptions;
import org.apache.hudi.io.FileGroupReaderBasedMergeHandle;
import org.apache.hudi.index.HoodieIndex;
import org.apache.hudi.table.HoodieTable;
import org.apache.hudi.table.marker.DirectWriteMarkers;
import org.apache.hudi.table.marker.TimelineServerBasedWriteMarkers;
import org.apache.hudi.table.marker.WriteMarkersFactory;
import org.apache.hudi.util.FlinkWriteClients;
import org.apache.hudi.util.StreamerUtil;

import org.apache.flink.configuration.Configuration;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;

import java.io.File;

import static org.hamcrest.CoreMatchers.is;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertTrue;

/**
 * Test cases for {@link org.apache.hudi.util.FlinkWriteClients}.
 */
public class TestFlinkWriteClients {
  @TempDir
  File tempFile;

  private Configuration conf;

  @BeforeEach
  public void before() throws Exception {
    this.conf = TestConfigurations.getDefaultConf(tempFile.getAbsolutePath());
  }

  @Test
  void testAutoSetupLockProvider() throws Exception {
    conf.set(FlinkOptions.METADATA_ENABLED, true);
    StreamerUtil.initTableIfNotExists(conf);
    HoodieWriteConfig writeConfig = FlinkWriteClients.getHoodieClientConfig(conf, false, false);
    assertThat(writeConfig.getLockProviderClass(), is(FileSystemBasedLockProvider.class.getName()));
    assertThat(writeConfig.getWriteConcurrencyMode(), is(WriteConcurrencyMode.SINGLE_WRITER));
  }

  @Test
  void testDefaultMetadataConfig() throws Exception {
    HoodieTableMetaClient metaClient = StreamerUtil.initTableIfNotExists(conf);
    HoodieWriteConfig writeConfig = FlinkWriteClients.getHoodieClientConfig(conf, false, false);
    assertTrue(writeConfig.isMetadataTableEnabled(), "MDT is by default enabled");
    assertFalse(writeConfig.isMetadataColumnStatsIndexEnabled(), "column_stats index is by default disabled");
    assertFalse(writeConfig.isSecondaryIndexEnabled(), "secondary index is by default disabled");
    assertFalse(writeConfig.isExpressionIndexEnabled(), "expression index is by default disabled");
    // create write client
    try (HoodieFlinkWriteClient writeClient = new HoodieFlinkWriteClient(HoodieFlinkEngineContext.DEFAULT, writeConfig)) {
      // init metadata table
      writeClient.initMetadataTable();
      // reload the table config been updated by the metadata table
      metaClient.reloadTableConfig();
      assertThat(metaClient.getTableConfig().getMetadataPartitions().toString(), is("[files]"));
    }
  }

  @Test
  void testDefaultGlobalRecordIndexMinFileGroupCountForFlink() {
    conf.set(FlinkOptions.INDEX_TYPE, HoodieIndex.IndexType.GLOBAL_RECORD_LEVEL_INDEX.name());
    HoodieWriteConfig writeConfig = FlinkWriteClients.getHoodieClientConfig(conf, false, false);
    assertEquals(8, writeConfig.getGlobalRecordLevelIndexMinFileGroupCount());
  }

  @Test
  void testUserConfiguredGlobalRecordIndexMinFileGroupCountIsNotOverridden() {
    conf.set(FlinkOptions.INDEX_TYPE, HoodieIndex.IndexType.GLOBAL_RECORD_LEVEL_INDEX.name());
    conf.setString(HoodieMetadataConfig.GLOBAL_RECORD_LEVEL_INDEX_MIN_FILE_GROUP_COUNT_PROP.key(), "12");
    HoodieWriteConfig writeConfig = FlinkWriteClients.getHoodieClientConfig(conf, false, false);
    assertEquals(12, writeConfig.getGlobalRecordLevelIndexMinFileGroupCount());
  }

  @ParameterizedTest
  @ValueSource(strings = {"", "DIRECT", "TIMELINE_SERVER_BASED"})
  void testMarkerType(String markerType) throws Exception {
    // create table
    StreamerUtil.initTableIfNotExists(conf);
    // This is expected to be used by the driver, the client can then send requests for files view.
    FlinkWriteClients.createWriteClient(conf);

    // do not set mark type to test the default behavior
    if (!markerType.isEmpty()) {
      conf.setString(HoodieWriteConfig.MARKERS_TYPE.key(), MarkerType.valueOf(markerType).name());
    }
    // This is expected to be used by writer client
    HoodieWriteConfig writeConfig = FlinkWriteClients.getHoodieClientConfig(conf, false, true);
    try (HoodieFlinkWriteClient writeClient = new HoodieFlinkWriteClient(HoodieFlinkEngineContext.DEFAULT, writeConfig)) {
      HoodieTable table = writeClient.getHoodieTable();
      String markerClass = WriteMarkersFactory.get(writeConfig.getMarkersType(), table, "001").getClass().getSimpleName();
      if (markerType.isEmpty() || markerType.equals("DIRECT")) {
        assertEquals(DirectWriteMarkers.class.getSimpleName(), markerClass);
      } else {
        assertEquals(TimelineServerBasedWriteMarkers.class.getSimpleName(), markerClass);
      }
    }
  }

  @ParameterizedTest
  @ValueSource(booleans = {true, false})
  void testRecordMergeConfigForEventTimeOrdering(boolean useLegacyConfig) throws Exception {
    if (useLegacyConfig) {
      conf.set(FlinkOptions.PAYLOAD_CLASS_NAME, EventTimeAvroPayload.class.getName());
    } else {
      conf.set(FlinkOptions.RECORD_MERGE_MODE, RecordMergeMode.EVENT_TIME_ORDERING.name());
    }
    HoodieTableMetaClient metaClient = StreamerUtil.initTableIfNotExists(conf);
    HoodieTableConfig tableConfig = metaClient.getTableConfig();

    assertThat(tableConfig.getRecordMergeMode(), is(RecordMergeMode.EVENT_TIME_ORDERING));
    assertNull(tableConfig.getRecordMergeStrategyId());
    assertThat(tableConfig.getPayloadClass(), is(DefaultHoodieRecordPayload.class.getName()));

    HoodieWriteConfig writeConfig = FlinkWriteClients.getHoodieClientConfig(conf, false, false);
    String mergerClasses = writeConfig.getString(HoodieWriteConfig.RECORD_MERGE_IMPL_CLASSES);
    assertThat(mergerClasses, is(EventTimeFlinkRecordMerger.class.getName()));
  }

  @ParameterizedTest
  @ValueSource(booleans = {true, false})
  void testRecordMergeConfigForCommitTimeOrdering(boolean useLegacyConfig) throws Exception {
    if (useLegacyConfig) {
      conf.set(FlinkOptions.PAYLOAD_CLASS_NAME, OverwriteWithLatestAvroPayload.class.getName());
    } else {
      conf.set(FlinkOptions.RECORD_MERGE_MODE, RecordMergeMode.COMMIT_TIME_ORDERING.name());
    }
    HoodieTableMetaClient metaClient = StreamerUtil.initTableIfNotExists(conf);
    HoodieTableConfig tableConfig = metaClient.getTableConfig();

    assertThat(tableConfig.getRecordMergeMode(), is(RecordMergeMode.COMMIT_TIME_ORDERING));
    assertNull(tableConfig.getRecordMergeStrategyId());
    assertThat(tableConfig.getPayloadClass(), is(OverwriteWithLatestAvroPayload.class.getName()));


    HoodieWriteConfig writeConfig = FlinkWriteClients.getHoodieClientConfig(conf, false, false);
    String mergerClasses = writeConfig.getString(HoodieWriteConfig.RECORD_MERGE_IMPL_CLASSES);
    assertThat(mergerClasses, is(CommitTimeFlinkRecordMerger.class.getName()));
  }

  @ParameterizedTest
  @ValueSource(ints = {1, 2, 3})
  void testRecordMergeConfigForPartialUpdate(int configOrdinal) throws Exception {
    if (configOrdinal == 1) {
      conf.set(FlinkOptions.PAYLOAD_CLASS_NAME, PartialUpdateAvroPayload.class.getName());
    } else if (configOrdinal == 2) {
      conf.set(FlinkOptions.RECORD_MERGER_IMPLS, PartialUpdateFlinkRecordMerger.class.getName());
    } else {
      conf.set(FlinkOptions.RECORD_MERGE_MODE, RecordMergeMode.CUSTOM.name());
      conf.set(FlinkOptions.RECORD_MERGER_IMPLS, PartialUpdateFlinkRecordMerger.class.getName());
    }
    conf.set(FlinkOptions.TABLE_TYPE, HoodieTableType.MERGE_ON_READ.name());
    HoodieTableMetaClient metaClient = StreamerUtil.initTableIfNotExists(conf);
    HoodieTableConfig tableConfig = metaClient.getTableConfig();

    assertThat(tableConfig.getRecordMergeMode(), is(RecordMergeMode.EVENT_TIME_ORDERING));
    assertNull(tableConfig.getRecordMergeStrategyId());
    assertThat(tableConfig.getPayloadClass(), is(DefaultHoodieRecordPayload.class.getName()));

    HoodieWriteConfig writeConfig = FlinkWriteClients.getHoodieClientConfig(conf, false, false);
    String mergerClasses = writeConfig.getString(HoodieWriteConfig.RECORD_MERGE_IMPL_CLASSES);
    assertThat(mergerClasses, is(PartialUpdateFlinkRecordMerger.class.getName()));
  }

  @Test
  void testWriteMergeHandleForPreV9Table() throws Exception {
    conf.set(FlinkOptions.WRITE_TABLE_VERSION, HoodieTableVersion.EIGHT.versionCode());
    StreamerUtil.initTableIfNotExists(conf);
    HoodieWriteConfig writeConfig = FlinkWriteClients.getHoodieClientConfig(conf, false, false);
    assertThat(writeConfig.getMergeHandleClassName(), is(FileGroupReaderBasedMergeHandle.class.getName()));
  }
}
