/*
 * 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 java.io.IOException;
import java.time.Duration;
import java.util.Collection;
import java.util.List;
import java.util.Locale;
import javax.annotation.Nullable;
import org.apache.flink.annotation.Internal;
import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.eventtime.TimestampAssigner;
import org.apache.flink.api.common.eventtime.TimestampAssignerSupplier;
import org.apache.flink.api.common.eventtime.Watermark;
import org.apache.flink.api.common.eventtime.WatermarkGenerator;
import org.apache.flink.api.common.eventtime.WatermarkGeneratorSupplier;
import org.apache.flink.api.common.eventtime.WatermarkOutput;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.operators.util.OperatorValidationUtils;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.connector.source.util.ratelimit.RateLimiterStrategy;
import org.apache.flink.api.dag.Transformation;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.DataStreamUtils;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.iceberg.flink.TableLoader;
import org.apache.iceberg.flink.maintenance.operator.LockRemover;
import org.apache.iceberg.flink.maintenance.operator.LockRemoverOperatorFactory;
import org.apache.iceberg.flink.maintenance.operator.MonitorSource;
import org.apache.iceberg.flink.maintenance.operator.TableChange;
import org.apache.iceberg.flink.maintenance.operator.TriggerEvaluator;
import org.apache.iceberg.flink.maintenance.operator.TriggerManager;
import org.apache.iceberg.flink.maintenance.operator.TriggerManagerOperatorFactory;
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.Lists;

/** Creates the table maintenance graph. */
public class TableMaintenance {
  static final String SOURCE_OPERATOR_NAME_PREFIX = "Monitor source for ";
  static final String TRIGGER_MANAGER_OPERATOR_NAME = "Trigger manager";
  static final String WATERMARK_ASSIGNER_OPERATOR_NAME = "Watermark Assigner";
  static final String FILTER_OPERATOR_NAME_PREFIX = "Filter ";
  static final String LOCK_REMOVER_OPERATOR_NAME = "Lock remover";

  static final long RATE_LIMIT_SECOND_DEFAULT = 60;
  static final long LOCK_CHECK_DELAY_SECOND_DEFAULT = 30;
  static final int MAX_READ_BACK_DEFAULT = 100;

  private TableMaintenance() {}

  /**
   * Use when the change stream is already provided, like in the {@link
   * IcebergSink#addPostCommitTopology(DataStream)}.
   *
   * @param changeStream the table changes
   * @param tableLoader used for accessing the table
   * @param lockFactory used for preventing concurrent task runs, if null, use coordination lock.
   * @return builder for the maintenance stream
   * @deprecated since 1.12.0, will be removed in 2.0.0. Use {@link #forChangeStream(DataStream,
   *     TableLoader)} instead.
   */
  @Deprecated
  @Internal
  public static Builder forChangeStream(
      DataStream<TableChange> changeStream,
      TableLoader tableLoader,
      @Nullable TriggerLockFactory lockFactory) {
    Preconditions.checkNotNull(changeStream, "The change stream should not be null");
    Preconditions.checkNotNull(tableLoader, "TableLoader should not be null");

    return new Builder(null, changeStream, tableLoader, lockFactory);
  }

  /**
   * Use when the change stream is already provided, like in the {@link
   * IcebergSink#addPostCommitTopology(DataStream)}.
   *
   * @param changeStream the table changes
   * @param tableLoader used for accessing the table
   * @return builder for the maintenance stream
   */
  @Internal
  public static Builder forChangeStream(
      DataStream<TableChange> changeStream, TableLoader tableLoader) {
    Preconditions.checkNotNull(changeStream, "The change stream should not be null");
    Preconditions.checkNotNull(tableLoader, "TableLoader should not be null");

    return new Builder(null, changeStream, tableLoader, null);
  }

  /**
   * Use this for standalone maintenance job. It creates a monitor source that detect table changes
   * and build the maintenance pipelines afterwards.
   *
   * @param env used to register the monitor source
   * @param tableLoader used for accessing the table
   * @param lockFactory used for preventing concurrent task runs. If null, use coordination lock.
   * @return builder for the maintenance stream
   * @deprecated since 1.12.0, will be removed in 2.0.0. Use {@link
   *     #forTable(StreamExecutionEnvironment, TableLoader)} instead.
   */
  @Deprecated
  public static Builder forTable(
      StreamExecutionEnvironment env,
      TableLoader tableLoader,
      @Nullable TriggerLockFactory lockFactory) {
    Preconditions.checkNotNull(env, "StreamExecutionEnvironment should not be null");
    Preconditions.checkNotNull(tableLoader, "TableLoader should not be null");

    return new Builder(env, null, tableLoader, lockFactory);
  }

  /**
   * Use this for standalone maintenance job. It creates a monitor source that detect table changes
   * and build the maintenance pipelines afterwards. But use coordination lock default.
   *
   * @param env used to register the monitor source
   * @param tableLoader used for accessing the table
   * @return builder for the maintenance stream
   */
  public static Builder forTable(StreamExecutionEnvironment env, TableLoader tableLoader) {
    Preconditions.checkNotNull(env, "StreamExecutionEnvironment should not be null");
    Preconditions.checkNotNull(tableLoader, "TableLoader should not be null");

    return new Builder(env, null, tableLoader, null);
  }

  public static class Builder {
    private final StreamExecutionEnvironment env;
    private final DataStream<TableChange> inputStream;
    private final TableLoader tableLoader;
    private final List<MaintenanceTaskBuilder<?>> taskBuilders;
    private final TriggerLockFactory lockFactory;

    private String uidSuffix = null;
    private String slotSharingGroup = null;
    private Duration rateLimit = Duration.ofSeconds(RATE_LIMIT_SECOND_DEFAULT);
    private Duration lockCheckDelay = Duration.ofSeconds(LOCK_CHECK_DELAY_SECOND_DEFAULT);
    private int parallelism = ExecutionConfig.PARALLELISM_DEFAULT;
    private int maxReadBack = MAX_READ_BACK_DEFAULT;

    private Builder(
        StreamExecutionEnvironment env,
        DataStream<TableChange> inputStream,
        TableLoader tableLoader,
        TriggerLockFactory lockFactory) {
      this.env = env;
      this.inputStream = inputStream;
      this.tableLoader = tableLoader;
      this.lockFactory = lockFactory;
      this.taskBuilders = Lists.newArrayListWithCapacity(4);
    }

    /**
     * The suffix used for the generated {@link Transformation}'s uid.
     *
     * @param newUidSuffix for the transformations
     */
    public Builder uidSuffix(String newUidSuffix) {
      this.uidSuffix = newUidSuffix;
      return this;
    }

    /**
     * The {@link SingleOutputStreamOperator#slotSharingGroup(String)} for all the operators of the
     * generated stream. Could be used to separate the resources used by this task.
     *
     * @param newSlotSharingGroup to be used for the operators
     */
    public Builder slotSharingGroup(String newSlotSharingGroup) {
      this.slotSharingGroup = newSlotSharingGroup;
      return this;
    }

    /**
     * Limits the firing frequency for the task triggers.
     *
     * @param newRateLimit firing frequency
     */
    public Builder rateLimit(Duration newRateLimit) {
      Preconditions.checkNotNull(newRateLimit, "Rate limit should not be null");
      Preconditions.checkArgument(
          newRateLimit.toMillis() > 0, "Rate limit should be greater than 0");
      this.rateLimit = newRateLimit;
      return this;
    }

    /**
     * Sets the delay for checking lock availability when a concurrent run is detected.
     *
     * @param newLockCheckDelay lock checking frequency
     */
    public Builder lockCheckDelay(Duration newLockCheckDelay) {
      this.lockCheckDelay = newLockCheckDelay;
      return this;
    }

    /**
     * Sets the default parallelism of maintenance tasks. Could be overwritten by the {@link
     * MaintenanceTaskBuilder#parallelism(int)}.
     *
     * @param newParallelism task parallelism
     */
    public Builder parallelism(int newParallelism) {
      OperatorValidationUtils.validateParallelism(newParallelism);
      this.parallelism = newParallelism;
      return this;
    }

    /**
     * Maximum number of snapshots checked when started with an embedded {@link MonitorSource} at
     * the first time. Only available when the {@link
     * TableMaintenance#forTable(StreamExecutionEnvironment, TableLoader, TriggerLockFactory)} is
     * used.
     *
     * @param newMaxReadBack snapshots to consider when initializing
     */
    public Builder maxReadBack(int newMaxReadBack) {
      Preconditions.checkArgument(
          inputStream == null, "Can't set maxReadBack when change stream is provided");
      this.maxReadBack = newMaxReadBack;
      return this;
    }

    /**
     * Adds a specific task with the given schedule.
     *
     * @param task to add
     */
    public Builder add(MaintenanceTaskBuilder<?> task) {
      taskBuilders.add(task);
      return this;
    }

    /**
     * Adds multiple tasks with the given schedules.
     *
     * @param tasks to add
     */
    public Builder add(Collection<MaintenanceTaskBuilder<?>> tasks) {
      Preconditions.checkNotNull(tasks, "Tasks collection should not be null");
      taskBuilders.addAll(tasks);
      return this;
    }

    /** Builds the task graph for the maintenance tasks. */
    public void append() throws IOException {
      Preconditions.checkArgument(!taskBuilders.isEmpty(), "Provide at least one task");

      List<String> taskNames = Lists.newArrayListWithCapacity(taskBuilders.size());
      List<TriggerEvaluator> evaluators = Lists.newArrayListWithCapacity(taskBuilders.size());
      for (int i = 0; i < taskBuilders.size(); ++i) {
        taskNames.add(nameFor(taskBuilders.get(i), i));
        evaluators.add(taskBuilders.get(i).evaluator());
      }

      try (TableLoader loader = tableLoader.clone()) {
        loader.open();
        String tableName = loader.loadTable().name();
        if (uidSuffix == null) {
          this.uidSuffix = "TableMaintenance-" + tableName;
        }

        DataStream<Trigger> triggers;
        if (lockFactory == null) {
          triggers =
              setSlotSharingGroup(
                  DataStreamUtils.reinterpretAsKeyedStream(
                          changeStream(tableName, loader), unused -> true)
                      .transform(
                          TRIGGER_MANAGER_OPERATOR_NAME,
                          TypeInformation.of(Trigger.class),
                          new TriggerManagerOperatorFactory(
                              tableName,
                              taskNames,
                              evaluators,
                              rateLimit.toMillis(),
                              lockCheckDelay.toMillis()))
                      .uid(TRIGGER_MANAGER_OPERATOR_NAME + uidSuffix)
                      .forceNonParallel());
        } else {
          triggers =
              setSlotSharingGroup(
                  DataStreamUtils.reinterpretAsKeyedStream(
                          changeStream(tableName, loader), unused -> true)
                      .process(
                          new TriggerManager(
                              loader,
                              lockFactory,
                              taskNames,
                              evaluators,
                              rateLimit.toMillis(),
                              lockCheckDelay.toMillis()))
                      .name(TRIGGER_MANAGER_OPERATOR_NAME)
                      .uid(TRIGGER_MANAGER_OPERATOR_NAME + uidSuffix)
                      .forceNonParallel());
        }

        triggers =
            setSlotSharingGroup(
                triggers
                    .assignTimestampsAndWatermarks(new PunctuatedWatermarkStrategy())
                    .name(WATERMARK_ASSIGNER_OPERATOR_NAME)
                    .uid(WATERMARK_ASSIGNER_OPERATOR_NAME + uidSuffix)
                    .forceNonParallel());

        // Add the specific tasks
        DataStream<TaskResult> unioned = null;
        for (int i = 0; i < taskBuilders.size(); ++i) {
          int taskIndex = i;
          DataStream<Trigger> filtered =
              setSlotSharingGroup(
                  triggers
                      .filter(t -> t.taskId() != null && t.taskId() == taskIndex)
                      .name(FILTER_OPERATOR_NAME_PREFIX + taskIndex)
                      .forceNonParallel()
                      .uid(FILTER_OPERATOR_NAME_PREFIX + taskIndex + "-" + uidSuffix));
          MaintenanceTaskBuilder<?> builder = taskBuilders.get(taskIndex);
          DataStream<TaskResult> result =
              builder.append(
                  filtered,
                  tableName,
                  taskNames.get(taskIndex),
                  taskIndex,
                  loader,
                  uidSuffix,
                  slotSharingGroup,
                  parallelism);
          if (unioned == null) {
            unioned = result;
          } else {
            unioned = unioned.union(result);
          }
        }

        // Add the LockRemover to the end
        if (lockFactory == null) {
          setSlotSharingGroup(
              unioned
                  .transform(
                      LOCK_REMOVER_OPERATOR_NAME,
                      TypeInformation.of(Void.class),
                      new LockRemoverOperatorFactory(tableName, taskNames))
                  .uid("lock-remover-" + uidSuffix)
                  .forceNonParallel());
        } else {
          setSlotSharingGroup(
              unioned
                  .transform(
                      LOCK_REMOVER_OPERATOR_NAME,
                      TypeInformation.of(Void.class),
                      new LockRemover(tableName, lockFactory, taskNames))
                  .forceNonParallel()
                  .uid("lock-remover-" + uidSuffix));
        }
      }
    }

    private DataStream<TableChange> changeStream(String tableName, TableLoader loader) {
      if (inputStream == null) {
        // Create a monitor source to provide the TableChange stream
        MonitorSource source =
            new MonitorSource(
                loader,
                RateLimiterStrategy.perSecond(monitorRatePerSecond(rateLimit.toMillis())),
                maxReadBack);
        return setSlotSharingGroup(
            env.fromSource(
                    source,
                    WatermarkStrategy.noWatermarks(),
                    SOURCE_OPERATOR_NAME_PREFIX + tableName)
                .uid(SOURCE_OPERATOR_NAME_PREFIX + uidSuffix)
                .forceNonParallel());
      } else {
        return inputStream.global();
      }
    }

    private static String nameFor(MaintenanceTaskBuilder<?> streamBuilder, int taskIndex) {
      return String.format(Locale.ROOT, "%s [%d]", streamBuilder.maintenanceTaskName(), taskIndex);
    }

    /**
     * Monitor poll rate per rate-limit interval, in checks/second. We compute from millis instead
     * of seconds, otherwise sub-second intervals could be truncated to 0, yielding an infinite rate
     * which busy-loops the source.
     */
    @VisibleForTesting
    static double monitorRatePerSecond(long rateLimitMillis) {
      return 1000.0 / rateLimitMillis;
    }

    private <T> SingleOutputStreamOperator<T> setSlotSharingGroup(
        SingleOutputStreamOperator<T> operator) {
      return slotSharingGroup == null ? operator : operator.slotSharingGroup(slotSharingGroup);
    }
  }

  @Internal
  public static class PunctuatedWatermarkStrategy implements WatermarkStrategy<Trigger> {
    @Override
    public WatermarkGenerator<Trigger> createWatermarkGenerator(
        WatermarkGeneratorSupplier.Context context) {
      return new WatermarkGenerator<>() {
        @Override
        public void onEvent(Trigger event, long eventTimestamp, WatermarkOutput output) {
          output.emitWatermark(new Watermark(event.timestamp()));
        }

        @Override
        public void onPeriodicEmit(WatermarkOutput output) {
          // No periodic watermarks
        }
      };
    }

    @Override
    public TimestampAssigner<Trigger> createTimestampAssigner(
        TimestampAssignerSupplier.Context context) {
      return (element, unused) -> element.timestamp();
    }
  }
}
