/*
 * 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.beam.runners.flink.translation.wrappers.streaming.io.source.unbounded;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import javax.annotation.Nonnull;
import org.apache.beam.runners.flink.FlinkPipelineOptions;
import org.apache.beam.runners.flink.translation.wrappers.streaming.io.source.FlinkSourceReaderBase;
import org.apache.beam.runners.flink.translation.wrappers.streaming.io.source.FlinkSourceSplit;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.io.Source;
import org.apache.beam.sdk.io.UnboundedSource;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
import org.apache.beam.sdk.values.ValueWithRecordId;
import org.apache.beam.sdk.values.WindowedValue;
import org.apache.beam.sdk.values.WindowedValues;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
import org.apache.flink.api.common.eventtime.Watermark;
import org.apache.flink.api.connector.source.ReaderOutput;
import org.apache.flink.api.connector.source.SourceOutput;
import org.apache.flink.api.connector.source.SourceReaderContext;
import org.apache.flink.core.io.InputStatus;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.joda.time.Instant;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * A Flink {@link org.apache.flink.api.connector.source.SourceReader SourceReader} implementation
 * that reads from the assigned {@link
 * org.apache.beam.runners.flink.translation.wrappers.streaming.io.source.FlinkSourceSplit
 * FlinkSourceSplits} by using Beam {@link org.apache.beam.sdk.io.UnboundedSource.UnboundedReader
 * UnboundedReaders}.
 *
 * <p>This reader consumes all the assigned source splits concurrently.
 *
 * @param <T> the output element type of the encapsulated Beam {@link
 *     org.apache.beam.sdk.io.UnboundedSource.UnboundedReader UnboundedReader}.
 */
public class FlinkUnboundedSourceReader<T>
    extends FlinkSourceReaderBase<T, WindowedValue<ValueWithRecordId<T>>> {
  private static final Logger LOG = LoggerFactory.getLogger(FlinkUnboundedSourceReader.class);
  // This name is defined in FLIP-33.
  @VisibleForTesting protected static final String PENDING_BYTES_METRIC_NAME = "pendingBytes";
  private static final long SLEEP_ON_IDLE_MS = 50L;
  private static final long MIN_WATERMARK_EMIT_INTERVAL_MS = 10L;
  private final AtomicReference<@Nullable CompletableFuture<Void>> dataAvailableFutureRef =
      new AtomicReference<>();
  private final List<ReaderAndOutput> readers = new ArrayList<>();
  private int currentReaderIndex = 0;
  private volatile boolean shouldEmitWatermark;
  private final LinkedHashMap<Long, List<FlinkSourceSplit<T>>> pendingCheckpoints =
      new LinkedHashMap<>();

  public FlinkUnboundedSourceReader(
      String stepName,
      SourceReaderContext context,
      PipelineOptions pipelineOptions,
      @Nullable Function<WindowedValue<ValueWithRecordId<T>>, Long> timestampExtractor) {
    super(stepName, context, pipelineOptions, timestampExtractor);
  }

  @VisibleForTesting
  protected FlinkUnboundedSourceReader(
      String stepName,
      SourceReaderContext context,
      PipelineOptions pipelineOptions,
      ScheduledExecutorService executor,
      @Nullable Function<WindowedValue<ValueWithRecordId<T>>, Long> timestampExtractor) {
    super(stepName, executor, context, pipelineOptions, timestampExtractor);
  }

  @Override
  protected void addSplitsToUnfinishedForCheckpoint(
      long checkpointId, List<FlinkSourceSplit<T>> flinkSourceSplits) {

    pendingCheckpoints.put(checkpointId, flinkSourceSplits);
  }

  @Override
  public void notifyCheckpointComplete(long checkpointId) throws Exception {
    super.notifyCheckpointComplete(checkpointId);
    List<Long> finalized = new ArrayList<>();
    for (Map.Entry<Long, List<FlinkSourceSplit<T>>> e : pendingCheckpoints.entrySet()) {
      if (e.getKey() <= checkpointId) {
        for (FlinkSourceSplit<T> s : e.getValue()) {
          finalizeSourceSplit(s.getCheckpointMark());
        }
        finalized.add(e.getKey());
      }
    }
    finalized.forEach(pendingCheckpoints::remove);
  }

  @Override
  public void start() {
    createPendingBytesGauge(context);
    Long watermarkInterval =
        pipelineOptions.as(FlinkPipelineOptions.class).getAutoWatermarkInterval();
    if (watermarkInterval == null) {
      watermarkInterval =
          (pipelineOptions.as(FlinkPipelineOptions.class).getMaxBundleTimeMills()) / 5L;
      watermarkInterval =
          (watermarkInterval > MIN_WATERMARK_EMIT_INTERVAL_MS)
              ? watermarkInterval
              : MIN_WATERMARK_EMIT_INTERVAL_MS;
      LOG.warn(
          "AutoWatermarkInterval is not set, watermarks will be emitted at a default interval of {} ms",
          watermarkInterval);
    }
    scheduleTaskAtFixedRate(
        () -> {
          // Set the watermark emission flag first.
          shouldEmitWatermark = true;
          // Wake up the main thread if necessary.
          CompletableFuture<Void> f = dataAvailableFutureRef.get();
          if (f != null) {
            f.complete(null);
          }
        },
        watermarkInterval,
        watermarkInterval);
  }

  @Override
  public InputStatus pollNext(ReaderOutput<WindowedValue<ValueWithRecordId<T>>> output)
      throws Exception {
    checkExceptionAndMaybeThrow();
    maybeEmitWatermark();
    maybeCreateReaderForNewSplits();

    ReaderAndOutput reader = nextReaderWithData(output);
    if (reader != null) {
      emitRecord(reader, output);
      return InputStatus.MORE_AVAILABLE;
    } else if (noMoreSplits() && isEndOfAllReaders()) {
      LOG.info("No more splits and no reader available. Terminating consumption.");
      return InputStatus.END_OF_INPUT;
    } else {
      LOG.trace("No data available for now.");
      return InputStatus.NOTHING_AVAILABLE;
    }
  }

  private boolean isEndOfAllReaders() {
    return allReaders().values().stream()
        .allMatch(
            r ->
                asUnbounded(r.reader).getWatermark().getMillis()
                    >= BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis());
  }

  /**
   * Check whether there are data available from alive readers. If not, set a future and wait for
   * the periodically running wake-up task to complete that future when the check interval passes.
   * This method is only called by the main thread, which is the only thread writing to the future
   * ref. Note that for UnboundedSource, because the splits never finishes, there are always alive
   * readers after the first split assigment. Hence, the return value of {@link
   * FlinkSourceReaderBase#isAvailable()} will effectively be determined by this method after the
   * first split assignment.
   */
  @Override
  protected CompletableFuture<Void> isAvailableForAliveReaders() {
    CompletableFuture<Void> future = dataAvailableFutureRef.get();
    if (future == null) {
      CompletableFuture<Void> newFuture = new CompletableFuture<>();
      // Need to set the future first to avoid the race condition of missing the watermark emission
      // notification.
      dataAvailableFutureRef.set(newFuture);
      if (shouldEmitWatermark || hasException()) {
        // There are exception after we set the new future,
        // immediately complete the future and return.
        dataAvailableFutureRef.set(null);
        newFuture.complete(null);
      } else {
        LOG.debug("There is no data available, scheduling the idle reader checker.");
        scheduleTask(
            () -> {
              CompletableFuture<Void> f = dataAvailableFutureRef.get();
              if (f != null) {
                f.complete(null);
              }
            },
            SLEEP_ON_IDLE_MS);
      }
      return newFuture;
    } else if (future.isDone()) {
      // The previous future is completed, just use it and reset the future ref.
      dataAvailableFutureRef.compareAndSet(future, null);
      return future;
    } else {
      // The previous future has not been completed, just use it.
      return future;
    }
  }

  @Override
  protected FlinkSourceSplit<T> getReaderCheckpoint(int splitId, ReaderAndOutput readerAndOutput) {
    // The checkpoint for unbounded sources is fine granular.
    UnboundedSource.UnboundedReader<T> reader =
        (UnboundedSource.UnboundedReader<T>) readerAndOutput.reader;
    UnboundedSource.CheckpointMark checkpointMark = reader.getCheckpointMark();
    @SuppressWarnings("unchecked")
    Coder<UnboundedSource.CheckpointMark> coder =
        (Coder<UnboundedSource.CheckpointMark>) reader.getCurrentSource().getCheckpointMarkCoder();
    byte[] checkpointState = encodeCheckpointMark(coder, checkpointMark);

    return new FlinkSourceSplit<>(
        splitId, readerAndOutput.reader.getCurrentSource(), checkpointState, checkpointMark);
  }

  @Override
  protected Source.Reader<T> createReader(@Nonnull FlinkSourceSplit<T> sourceSplit)
      throws IOException {
    Source<T> beamSource = sourceSplit.getBeamSplitSource();
    return createUnboundedSourceReader(beamSource, sourceSplit.getSplitState());
  }

  // -------------- private helper methods ----------------

  private void emitRecord(
      ReaderAndOutput readerAndOutput, ReaderOutput<WindowedValue<ValueWithRecordId<T>>> output) {
    UnboundedSource.UnboundedReader<T> reader = asUnbounded(readerAndOutput.reader);
    T item = reader.getCurrent();
    byte[] recordId = reader.getCurrentRecordId();
    Instant timestamp = reader.getCurrentTimestamp();

    WindowedValue<ValueWithRecordId<T>> windowedValue =
        WindowedValues.of(
            new ValueWithRecordId<>(item, recordId),
            timestamp,
            GlobalWindow.INSTANCE,
            PaneInfo.NO_FIRING);
    LOG.trace("Emitting record: {}", windowedValue);
    if (timestampExtractor == null) {
      readerAndOutput.getAndMaybeCreateSplitOutput(output).collect(windowedValue);
    } else {
      readerAndOutput
          .getAndMaybeCreateSplitOutput(output)
          .collect(windowedValue, timestampExtractor.apply(windowedValue));
    }
    numRecordsInCounter.inc();
  }

  private void maybeEmitWatermark() {
    // Here we rely on the Flink source watermark multiplexer to combine the per-split watermark.
    // The runnable may emit more than one watermark when it runs.
    if (shouldEmitWatermark) {
      allReaders()
          .values()
          .forEach(
              readerAndOutput -> {
                SourceOutput<?> sourceOutput = readerAndOutput.sourceOutput();
                if (sourceOutput != null) {
                  long watermark = asUnbounded(readerAndOutput.reader).getWatermark().getMillis();
                  sourceOutput.emitWatermark(new Watermark(watermark));
                }
              });
      shouldEmitWatermark = false;
    }
  }

  private void maybeCreateReaderForNewSplits() throws Exception {
    while (!sourceSplits().isEmpty()) {
      Optional<ReaderAndOutput> readerAndOutputOpt = createAndTrackNextReader();
      if (readerAndOutputOpt.isPresent()) {
        readers.add(readerAndOutputOpt.get());
      } else {
        // Null splitId is only possible when exception occurs, just check exception to throw it.
        checkExceptionAndMaybeThrow();
      }
    }
  }

  private @Nullable ReaderAndOutput nextReaderWithData(
      ReaderOutput<WindowedValue<ValueWithRecordId<T>>> output) throws IOException {

    int numReaders = readers.size();
    for (int i = 0; i < numReaders; i++) {
      ReaderAndOutput readerAndOutput = readers.get(currentReaderIndex);
      currentReaderIndex = (currentReaderIndex + 1) % numReaders;
      if (readerAndOutput.startOrAdvance(output)) {
        return readerAndOutput;
      }
    }
    return null;
  }

  private static <T> UnboundedSource.UnboundedReader<T> asUnbounded(Source.Reader<T> reader) {
    return (UnboundedSource.UnboundedReader<T>) reader;
  }

  private void createPendingBytesGauge(SourceReaderContext context) {
    // TODO: Replace with SourceReaderContest.metricGroup().setPendingBytesGauge() after Flink 1.14
    // and above.
    context
        .metricGroup()
        .gauge(
            PENDING_BYTES_METRIC_NAME,
            () -> {
              long pendingBytes = -1L;
              for (FlinkSourceReaderBase<?, ?>.ReaderAndOutput readerAndOutput :
                  allReaders().values()) {
                long pendingBytesForReader =
                    asUnbounded(readerAndOutput.reader).getSplitBacklogBytes();
                if (pendingBytesForReader != -1L) {
                  pendingBytes =
                      pendingBytes == -1L
                          ? pendingBytesForReader
                          : pendingBytes + pendingBytesForReader;
                }
              }
              return pendingBytes;
            });
  }

  private <CheckpointMarkT extends UnboundedSource.CheckpointMark> byte[] encodeCheckpointMark(
      Coder<CheckpointMarkT> coder, CheckpointMarkT checkpointMark) {

    try (ByteArrayOutputStream baos = new ByteArrayOutputStream()) {
      coder.encode(checkpointMark, baos);
      return baos.toByteArray();
    } catch (IOException ioe) {
      throw new RuntimeException("Failed to encode checkpoint mark.", ioe);
    }
  }

  private <CheckpointMarkT extends UnboundedSource.CheckpointMark>
      Source.Reader<T> createUnboundedSourceReader(
          Source<T> beamSource, byte @Nullable [] splitState) throws IOException {
    UnboundedSource<T, CheckpointMarkT> unboundedSource =
        (UnboundedSource<T, CheckpointMarkT>) beamSource;
    Coder<CheckpointMarkT> coder = unboundedSource.getCheckpointMarkCoder();
    if (splitState == null) {
      return unboundedSource.createReader(pipelineOptions, null);
    } else {
      try (ByteArrayInputStream bais = new ByteArrayInputStream(splitState)) {
        return unboundedSource.createReader(pipelineOptions, coder.decode(bais));
      }
    }
  }

  private void finalizeSourceSplit(UnboundedSource.@Nullable CheckpointMark mark)
      throws IOException {
    if (mark != null) {
      mark.finalizeCheckpoint();
    }
  }
}
