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

import java.util.Map;
import java.util.Optional;
import java.util.UUID;
import org.apache.flink.annotation.Experimental;
import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.connector.sink2.Committer;
import org.apache.flink.api.connector.sink2.CommitterInitContext;
import org.apache.flink.api.connector.sink2.Sink;
import org.apache.flink.api.connector.sink2.SinkWriter;
import org.apache.flink.api.connector.sink2.SupportsCommitter;
import org.apache.flink.api.connector.sink2.WriterInitContext;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.configuration.ReadableConfig;
import org.apache.flink.core.io.SimpleVersionedSerializer;
import org.apache.flink.streaming.api.connector.sink2.CommittableMessage;
import org.apache.flink.streaming.api.connector.sink2.CommittableMessageTypeInfo;
import org.apache.flink.streaming.api.connector.sink2.CommittableSummary;
import org.apache.flink.streaming.api.connector.sink2.CommittableWithLineage;
import org.apache.flink.streaming.api.connector.sink2.SupportsPostCommitTopology;
import org.apache.flink.streaming.api.connector.sink2.SupportsPreCommitTopology;
import org.apache.flink.streaming.api.connector.sink2.SupportsPreWriteTopology;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.DataStreamSink;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.runtime.operators.sink.SinkWriterOperatorFactory;
import org.apache.flink.table.data.RowData;
import org.apache.flink.util.OutputTag;
import org.apache.iceberg.Table;
import org.apache.iceberg.flink.CatalogLoader;
import org.apache.iceberg.flink.FlinkWriteConf;
import org.apache.iceberg.flink.FlinkWriteOptions;
import org.apache.iceberg.flink.sink.IcebergSink;
import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
import org.apache.iceberg.relocated.com.google.common.collect.Maps;

/**
 * Dynamic version of the IcebergSink which supports:
 *
 * <ul>
 *   <li>Writing to any number of tables (No more 1:1 sink/topic relationship).
 *   <li>Creating and updating tables based on the user-supplied routing.
 *   <li>Updating the schema and partition spec of tables based on the user-supplied specification.
 * </ul>
 */
@Experimental
public class DynamicIcebergSink
    implements Sink<DynamicRecordInternal>,
        SupportsPreWriteTopology<DynamicRecordInternal>,
        SupportsCommitter<DynamicCommittable>,
        SupportsPreCommitTopology<DynamicWriteResult, DynamicCommittable>,
        SupportsPostCommitTopology<DynamicCommittable> {

  private final CatalogLoader catalogLoader;
  private final Map<String, String> snapshotProperties;
  private final String uidPrefix;
  private final String sinkId;
  private final Map<String, String> writeProperties;
  private final Configuration flinkConfig;
  private final int cacheMaximumSize;

  // Set by the builder before sinkTo() — forward writer results to union into pre-commit topology
  private final transient DataStream<CommittableMessage<DynamicWriteResult>> forwardWriteResults;

  DynamicIcebergSink(
      CatalogLoader catalogLoader,
      Map<String, String> snapshotProperties,
      String uidPrefix,
      Map<String, String> writeProperties,
      Configuration flinkConfig,
      int cacheMaximumSize,
      DataStream<CommittableMessage<DynamicWriteResult>> forwardWriteResults) {
    this.catalogLoader = catalogLoader;
    this.snapshotProperties = snapshotProperties;
    this.uidPrefix = uidPrefix;
    this.writeProperties = writeProperties;
    this.flinkConfig = flinkConfig;
    this.cacheMaximumSize = cacheMaximumSize;
    // We generate a random UUID every time when a sink is created.
    // This is used to separate files generated by different sinks writing the same table.
    // Also used to generate the aggregator operator name
    this.sinkId = UUID.randomUUID().toString();
    this.forwardWriteResults = forwardWriteResults;
  }

  @Override
  public SinkWriter<DynamicRecordInternal> createWriter(WriterInitContext context) {
    return new DynamicWriter(
        catalogLoader.loadCatalog(),
        writeProperties,
        flinkConfig,
        cacheMaximumSize,
        new DynamicWriterMetrics(context.metricGroup()),
        context.getTaskInfo().getIndexOfThisSubtask(),
        context.getTaskInfo().getAttemptNumber());
  }

  @Override
  public Committer<DynamicCommittable> createCommitter(CommitterInitContext context) {
    FlinkWriteConf flinkWriteConf = new FlinkWriteConf(writeProperties, flinkConfig);
    DynamicCommitterMetrics metrics = new DynamicCommitterMetrics(context.metricGroup());
    return new DynamicCommitter(
        catalogLoader.loadCatalog(),
        snapshotProperties,
        flinkWriteConf.overwriteMode(),
        flinkWriteConf.workerPoolSize(),
        sinkId,
        metrics);
  }

  @Override
  public SimpleVersionedSerializer<DynamicCommittable> getCommittableSerializer() {
    return new DynamicCommittableSerializer();
  }

  @Override
  public void addPostCommitTopology(
      DataStream<CommittableMessage<DynamicCommittable>> committables) {}

  @Override
  public DataStream<DynamicRecordInternal> addPreWriteTopology(
      DataStream<DynamicRecordInternal> inputDataStream) {
    return distributeDataStream(inputDataStream);
  }

  @Override
  public DataStream<CommittableMessage<DynamicCommittable>> addPreCommitTopology(
      DataStream<CommittableMessage<DynamicWriteResult>> writeResults) {
    TypeInformation<CommittableMessage<DynamicCommittable>> typeInformation =
        CommittableMessageTypeInfo.of(this::getCommittableSerializer);

    // Union forward writer results with the shuffle writer results
    DataStream<CommittableMessage<DynamicWriteResult>> allResults =
        writeResults.union(forwardWriteResults);

    return allResults
        .keyBy(
            committable -> {
              if (committable instanceof CommittableSummary) {
                return "__summary";
              } else {
                CommittableWithLineage<DynamicWriteResult> result =
                    (CommittableWithLineage<DynamicWriteResult>) committable;
                return result.getCommittable().key().tableName();
              }
            })
        .transform(
            prefixIfNotNull(uidPrefix, sinkId + " Pre Commit"),
            typeInformation,
            new DynamicWriteResultAggregator(catalogLoader, cacheMaximumSize))
        .uid(prefixIfNotNull(uidPrefix, "-pre-commit-topology"));
  }

  @Override
  public SimpleVersionedSerializer<DynamicWriteResult> getWriteResultSerializer() {
    return new DynamicWriteResultSerializer();
  }

  /**
   * A lightweight Sink used with {@link SinkWriterOperatorFactory} for the forward write path.
   * Implements {@link SupportsCommitter} so that {@code SinkWriterOperator} emits committables
   * downstream. The committer is never called — committing is handled by the main sink.
   */
  @VisibleForTesting
  static class ForwardWriterSink
      implements Sink<DynamicRecordInternal>, SupportsCommitter<DynamicWriteResult> {

    private final CatalogLoader catalogLoader;
    private final Map<String, String> writeProperties;
    private final Configuration flinkConfig;
    private final int cacheMaximumSize;

    ForwardWriterSink(
        CatalogLoader catalogLoader,
        Map<String, String> writeProperties,
        Configuration flinkConfig,
        int cacheMaximumSize) {
      this.catalogLoader = catalogLoader;
      this.writeProperties = writeProperties;
      this.flinkConfig = flinkConfig;
      this.cacheMaximumSize = cacheMaximumSize;
    }

    @Override
    public SinkWriter<DynamicRecordInternal> createWriter(WriterInitContext context) {
      return new DynamicWriter(
          catalogLoader.loadCatalog(),
          writeProperties,
          flinkConfig,
          cacheMaximumSize,
          new DynamicWriterMetrics(context.metricGroup()),
          context.getTaskInfo().getIndexOfThisSubtask(),
          context.getTaskInfo().getAttemptNumber());
    }

    @Override
    public Committer<DynamicWriteResult> createCommitter(CommitterInitContext context) {
      throw new UnsupportedOperationException(
          "WriterSink is used only for writing; committing is handled by the main sink");
    }

    @Override
    public SimpleVersionedSerializer<DynamicWriteResult> getCommittableSerializer() {
      return new DynamicWriteResultSerializer();
    }
  }

  public static class Builder<T> {
    private DataStream<T> input;
    private DynamicRecordGenerator<T> generator;
    private CatalogLoader catalogLoader;
    private String uidPrefix = null;
    private final Map<String, String> writeOptions = Maps.newHashMap();
    private final Map<String, String> snapshotSummary = Maps.newHashMap();
    private ReadableConfig readableConfig = new Configuration();
    private TableCreator tableCreator = TableCreator.DEFAULT;

    Builder() {}

    public Builder<T> forInput(DataStream<T> inputStream) {
      this.input = inputStream;
      return this;
    }

    public Builder<T> generator(DynamicRecordGenerator<T> inputGenerator) {
      this.generator = inputGenerator;
      return this;
    }

    /**
     * The catalog loader is used for loading tables in {@link DynamicCommitter} lazily, we need
     * this loader because {@link Table} is not serializable and could not just use the loaded table
     * from Builder#table in the remote task manager.
     *
     * @param newCatalogLoader to load iceberg table inside tasks.
     * @return {@link Builder} to connect the iceberg table.
     */
    public Builder<T> catalogLoader(CatalogLoader newCatalogLoader) {
      this.catalogLoader = newCatalogLoader;
      return this;
    }

    /**
     * Set the write properties for IcebergSink. View the supported properties in {@link
     * FlinkWriteOptions}
     */
    public Builder<T> set(String property, String value) {
      writeOptions.put(property, value);
      return this;
    }

    /**
     * Set the write properties for IcebergSink. View the supported properties in {@link
     * FlinkWriteOptions}
     */
    public Builder<T> setAll(Map<String, String> properties) {
      writeOptions.putAll(properties);
      return this;
    }

    public Builder<T> overwrite(boolean newOverwrite) {
      writeOptions.put(FlinkWriteOptions.OVERWRITE_MODE.key(), Boolean.toString(newOverwrite));
      return this;
    }

    public Builder<T> flinkConf(ReadableConfig config) {
      this.readableConfig = config;
      return this;
    }

    /**
     * Logic to create a table. Allows setting custom table properties/location on a per-table
     * basis.
     */
    public Builder<T> tableCreator(TableCreator tableCreationFunction) {
      this.tableCreator = tableCreationFunction;
      return this;
    }

    /**
     * Configuring the write parallel number for iceberg stream writer.
     *
     * @param newWriteParallelism the number of parallel iceberg stream writer.
     * @return {@link DynamicIcebergSink.Builder} to connect the iceberg table.
     */
    public Builder<T> writeParallelism(int newWriteParallelism) {
      writeOptions.put(
          FlinkWriteOptions.WRITE_PARALLELISM.key(), Integer.toString(newWriteParallelism));
      return this;
    }

    /**
     * Set the slot sharing group for the generator (and the forward-writer chained to it) by name.
     * Register the slot sharing group with its resource spec on the {@code
     * StreamExecutionEnvironment} via {@code env.registerSlotSharingGroup(...)}. If unset, Flink
     * inherits the slot sharing group from the upstream operator.
     */
    public Builder<T> generatorSlotSharingGroup(String ssg) {
      writeOptions.put(FlinkDynamicSinkOptions.GENERATOR_SLOT_SHARING_GROUP.key(), ssg);
      return this;
    }

    /**
     * Set the slot sharing group for the shuffling sink (writer plus committer) by name. Register
     * the slot sharing group with its resource spec on the {@code StreamExecutionEnvironment} via
     * {@code env.registerSlotSharingGroup(...)}. If unset, Flink inherits the slot sharing group
     * from the upstream operator.
     */
    public Builder<T> shuffleSinkSlotSharingGroup(String ssg) {
      writeOptions.put(FlinkDynamicSinkOptions.SHUFFLE_SINK_SLOT_SHARING_GROUP.key(), ssg);
      return this;
    }

    /**
     * Set the uid prefix for IcebergSink operators. Note that IcebergSink internally consists of
     * multiple operators (like writer, committer, aggregator) Actual operator uid will be appended
     * with a suffix like "uidPrefix-writer".
     *
     * <p>If provided, this prefix is also applied to operator names.
     *
     * <p>Flink auto generates operator uid if not set explicitly. It is a recommended <a
     * href="https://ci.apache.org/projects/flink/flink-docs-master/docs/ops/production_ready/">
     * best-practice to set uid for all operators</a> before deploying to production. Flink has an
     * option to {@code pipeline.auto-generate-uid=false} to disable auto-generation and force
     * explicit setting of all operator uid.
     *
     * <p>Be careful with setting this for an existing job, because now we are changing the operator
     * uid from an auto-generated one to this new value. When deploying the change with a
     * checkpoint, Flink won't be able to restore the previous IcebergSink operator state (more
     * specifically the committer operator state). You need to use {@code --allowNonRestoredState}
     * to ignore the previous sink state. During restore IcebergSink state is used to check if last
     * commit was actually successful or not. {@code --allowNonRestoredState} can lead to data loss
     * if the Iceberg commit failed in the last completed checkpoint.
     *
     * @param newPrefix prefix for Flink sink operator uid and name
     * @return {@link Builder} to connect the iceberg table.
     */
    public Builder<T> uidPrefix(String newPrefix) {
      this.uidPrefix = newPrefix;
      return this;
    }

    public Builder<T> snapshotProperties(Map<String, String> properties) {
      snapshotSummary.putAll(properties);
      return this;
    }

    public Builder<T> setSnapshotProperty(String property, String value) {
      snapshotSummary.put(property, value);
      return this;
    }

    public Builder<T> toBranch(String branch) {
      writeOptions.put(FlinkWriteOptions.BRANCH.key(), branch);
      return this;
    }

    public Builder<T> immediateTableUpdate(boolean newImmediateUpdate) {
      writeOptions.put(
          FlinkDynamicSinkOptions.IMMEDIATE_TABLE_UPDATE.key(),
          Boolean.toString(newImmediateUpdate));
      return this;
    }

    /**
     * Dropping columns is disabled by default to prevent issues with late or out-of-order data, as
     * removed fields cannot be easily restored without data loss.
     *
     * <p>You can opt-in to allow dropping columns. Once a column has been dropped, it is
     * technically still possible to write data to that column because Iceberg maintains all past
     * table schemas. However, regular queries won't be able to reference the column. If the field
     * was to re-appear as part of a new schema, an entirely new column would be added, which apart
     * from the name, has nothing in common with the old column, i.e. queries for the new column
     * will never return data of the old column.
     */
    public Builder<T> dropUnusedColumns(boolean newDropUnusedColumns) {
      writeOptions.put(
          FlinkDynamicSinkOptions.DROP_UNUSED_COLUMNS.key(),
          Boolean.toString(newDropUnusedColumns));
      return this;
    }

    /** Maximum size of the caches used in Dynamic Sink for table data and serializers. */
    public Builder<T> cacheMaxSize(int maxSize) {
      writeOptions.put(FlinkDynamicSinkOptions.CACHE_MAX_SIZE.key(), Integer.toString(maxSize));
      return this;
    }

    /** Maximum interval for cache items renewals. */
    public Builder<T> cacheRefreshMs(long refreshMs) {
      writeOptions.put(FlinkDynamicSinkOptions.CACHE_REFRESH_MS.key(), Long.toString(refreshMs));
      return this;
    }

    /**
     * Maximum input {@link org.apache.iceberg.Schema} objects to cache per each Iceberg table. The
     * cache improves Dynamic Sink performance by storing {@link org.apache.iceberg.Schema}
     * comparison results.
     */
    public Builder<T> inputSchemasPerTableCacheMaxSize(int inputSchemasPerTableCacheMaxSize) {
      writeOptions.put(
          FlinkDynamicSinkOptions.INPUT_SCHEMAS_PER_TABLE_CACHE_MAX_SIZE.key(),
          Integer.toString(inputSchemasPerTableCacheMaxSize));
      return this;
    }

    /**
     * Set whether schema field name matching should be case-sensitive. The default is to match the
     * field names case-sensitive.
     */
    public Builder<T> caseSensitive(boolean newCaseSensitive) {
      writeOptions.put(
          FlinkDynamicSinkOptions.CASE_SENSITIVE.key(), Boolean.toString(newCaseSensitive));
      return this;
    }

    private String operatorName(String suffix) {
      return uidPrefix != null ? uidPrefix + "-" + suffix : suffix;
    }

    private DynamicIcebergSink build(
        SingleOutputStreamOperator<DynamicRecordInternal> converted,
        DynamicRecordInternalType sideOutputType) {

      Preconditions.checkArgument(
          generator != null, "Please use withGenerator() to convert the input DataStream.");
      Preconditions.checkNotNull(catalogLoader, "Catalog loader shouldn't be null");

      Configuration flinkConfig = fromReadableConfig();
      FlinkDynamicSinkConf flinkDynamicSinkConf =
          new FlinkDynamicSinkConf(writeOptions, flinkConfig);

      // Forward writer: chained with generator via forward edge, no data shuffle
      ForwardWriterSink forwardWriterSink =
          new ForwardWriterSink(
              catalogLoader, writeOptions, flinkConfig, flinkDynamicSinkConf.cacheMaxSize());
      TypeInformation<CommittableMessage<DynamicWriteResult>> writeResultTypeInfo =
          CommittableMessageTypeInfo.of(DynamicWriteResultSerializer::new);

      SingleOutputStreamOperator<CommittableMessage<DynamicWriteResult>> forwardWriteResults =
          converted
              .getSideOutput(
                  new OutputTag<>(DynamicRecordProcessor.DYNAMIC_FORWARD_STREAM, sideOutputType))
              .transform(
                  operatorName("Forward-Writer"),
                  writeResultTypeInfo,
                  new SinkWriterOperatorFactory<>(forwardWriterSink))
              .setParallelism(converted.getParallelism())
              .uid(prefixIfNotNull(uidPrefix, "-forward-writer"));

      String generatorSsg = flinkDynamicSinkConf.generatorSlotSharingGroup();
      if (generatorSsg != null) {
        forwardWriteResults.slotSharingGroup(generatorSsg);
      }

      // Inject forward write results into sink — they'll be unioned in addPreCommitTopology
      return instantiateSink(writeOptions, flinkConfig, forwardWriteResults);
    }

    @VisibleForTesting
    DynamicIcebergSink instantiateSink(
        Map<String, String> writeProperties,
        Configuration flinkWriteConf,
        DataStream<CommittableMessage<DynamicWriteResult>> forwardWriteResults) {
      FlinkDynamicSinkConf flinkDynamicSinkConf =
          new FlinkDynamicSinkConf(writeProperties, flinkWriteConf);
      return new DynamicIcebergSink(
          catalogLoader,
          snapshotSummary,
          uidPrefix,
          writeProperties,
          flinkWriteConf,
          flinkDynamicSinkConf.cacheMaxSize(),
          forwardWriteResults);
    }

    /**
     * Append the iceberg sink operators to write records to iceberg table.
     *
     * <p>The topology splits records by distribution mode:
     *
     * <ul>
     *   <li>Forward records ({@code null} distributionMode) go through a forward edge to a chained
     *       writer, avoiding any data shuffle.
     *   <li>Shuffle records (non-null distributionMode) go through the standard Sink2 pipeline with
     *       hash/round-robin distribution.
     * </ul>
     *
     * Both writers feed into a single shared pre-commit aggregator and committer, ensuring atomic
     * commits across both paths.
     *
     * @return {@link DataStreamSink} for sink.
     */
    public DataStreamSink<DynamicRecordInternal> append() {
      uidPrefix = Optional.ofNullable(uidPrefix).orElse("");

      FlinkDynamicSinkConf flinkDynamicSinkConf =
          new FlinkDynamicSinkConf(writeOptions, readableConfig);
      Configuration flinkConfig = fromReadableConfig();

      DynamicRecordInternalType type =
          new DynamicRecordInternalType(catalogLoader, false, flinkDynamicSinkConf.cacheMaxSize());
      DynamicRecordInternalType sideOutputType =
          new DynamicRecordInternalType(catalogLoader, true, flinkDynamicSinkConf.cacheMaxSize());

      SingleOutputStreamOperator<DynamicRecordInternal> converted =
          input
              .process(
                  new DynamicRecordProcessor<>(
                      generator,
                      catalogLoader,
                      tableCreator,
                      flinkDynamicSinkConf,
                      writeOptions,
                      flinkConfig))
              .setParallelism(input.getParallelism())
              .uid(prefixIfNotNull(uidPrefix, "-generator"))
              .name(operatorName("generator"))
              .returns(type);
      String generatorSsg = flinkDynamicSinkConf.generatorSlotSharingGroup();
      if (generatorSsg != null) {
        converted.slotSharingGroup(generatorSsg);
      }

      DynamicIcebergSink sink = build(converted, sideOutputType);

      // Shuffle path: table update side output + main output → sinkTo()
      DataStream<DynamicRecordInternal> shuffleInput =
          converted
              .getSideOutput(
                  new OutputTag<>(
                      DynamicRecordProcessor.DYNAMIC_TABLE_UPDATE_STREAM, sideOutputType))
              .keyBy((KeySelector<DynamicRecordInternal, String>) DynamicRecordInternal::tableName)
              .map(
                  new DynamicTableUpdateOperator(catalogLoader, tableCreator, flinkDynamicSinkConf))
              .uid(prefixIfNotNull(uidPrefix, "-updater"))
              .name(operatorName("Updater"))
              .returns(type)
              .union(converted);

      DataStreamSink<DynamicRecordInternal> result =
          shuffleInput
              .sinkTo(sink) // Forward write results are implicitly injected here
              .uid(prefixIfNotNull(uidPrefix, "-sink"));
      String shuffleSinkSsg = flinkDynamicSinkConf.shuffleSinkSlotSharingGroup();
      if (shuffleSinkSsg != null) {
        result.slotSharingGroup(shuffleSinkSsg);
      }

      FlinkWriteConf flinkWriteConf = new FlinkWriteConf(writeOptions, readableConfig);
      if (flinkWriteConf.writeParallelism() != null) {
        result.setParallelism(flinkWriteConf.writeParallelism());
      }

      return result;
    }

    private Configuration fromReadableConfig() {
      return readableConfig instanceof Configuration
          ? (Configuration) readableConfig
          : Configuration.fromMap(readableConfig.toMap());
    }
  }

  DataStream<DynamicRecordInternal> distributeDataStream(DataStream<DynamicRecordInternal> input) {
    return input.keyBy(DynamicRecordInternal::writerKey);
  }

  private static String prefixIfNotNull(String uidPrefix, String suffix) {
    return uidPrefix != null ? uidPrefix + "-" + suffix : suffix;
  }

  /**
   * Initialize a {@link IcebergSink.Builder} to export the data from input data stream with {@link
   * RowData}s into iceberg table.
   *
   * @param input the source input data stream with {@link RowData}s.
   * @return {@link IcebergSink.Builder} to connect the iceberg table.
   */
  public static <T> Builder<T> forInput(DataStream<T> input) {
    return new Builder<T>().forInput(input);
  }
}
