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

import static com.google.common.base.Preconditions.checkArgument;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.NoSuchElementException;
import org.apache.beam.sdk.io.range.OffsetRange;
import org.apache.beam.sdk.io.range.OffsetRangeTracker;
import org.apache.beam.sdk.io.range.RangeTracker;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * A {@link BoundedSource} that uses offsets to define starting and ending positions.
 *
 * <p>{@link OffsetBasedSource} is a common base class for all bounded sources where the input can
 * be represented as a single range, and an input can be efficiently processed in parallel by
 * splitting the range into a set of disjoint ranges whose union is the original range. This class
 * should be used for sources that can be cheaply read starting at any given offset.
 * {@link OffsetBasedSource} stores the range and implements splitting into bundles.
 *
 * <p>Extend {@link OffsetBasedSource} to implement your own offset-based custom source.
 * {@link FileBasedSource}, which is a subclass of this, adds additional functionality useful for
 * custom sources that are based on files. If possible implementors should start from
 * {@link FileBasedSource} instead of {@link OffsetBasedSource}.
 *
 * <p>Consult {@link RangeTracker} for important semantics common to all sources defined by a range
 * of positions of a certain type, including the semantics of split points
 * ({@link OffsetBasedReader#isAtSplitPoint}).
 *
 * @param <T> Type of records represented by the source.
 * @see BoundedSource
 * @see FileBasedSource
 * @see RangeTracker
 */
public abstract class OffsetBasedSource<T> extends BoundedSource<T> {
  private final long startOffset;
  private final long endOffset;
  private final long minBundleSize;

  /**
   * @param startOffset starting offset (inclusive) of the source. Must be non-negative.
   *
   * @param endOffset ending offset (exclusive) of the source. Use {@link Long#MAX_VALUE} to
   *        indicate that the entire source after {@code startOffset} should be read. Must be
   *        {@code > startOffset}.
   *
   * @param minBundleSize minimum bundle size in offset units that should be used when splitting the
   *                      source into sub-sources. This value may not be respected if the total
   *                      range of the source is smaller than the specified {@code minBundleSize}.
   *                      Must be non-negative.
   */
  public OffsetBasedSource(long startOffset, long endOffset, long minBundleSize) {
    this.startOffset = startOffset;
    this.endOffset = endOffset;
    this.minBundleSize = minBundleSize;
  }

  /**
   * Returns the starting offset of the source.
   */
  public long getStartOffset() {
    return startOffset;
  }

  /**
   * Returns the specified ending offset of the source. Any returned value greater than or equal to
   * {@link #getMaxEndOffset(PipelineOptions)} should be treated as
   * {@link #getMaxEndOffset(PipelineOptions)}.
   */
  public long getEndOffset() {
    return endOffset;
  }

  /**
   * Returns the minimum bundle size that should be used when splitting the source into sub-sources.
   * This value may not be respected if the total range of the source is smaller than the specified
   * {@code minBundleSize}.
   */
  public long getMinBundleSize() {
    return minBundleSize;
  }

  @Override
  public long getEstimatedSizeBytes(PipelineOptions options) throws Exception {
    long trueEndOffset = (endOffset == Long.MAX_VALUE) ? getMaxEndOffset(options) : endOffset;
    return getBytesPerOffset() * (trueEndOffset - getStartOffset());
  }

  @Override
  public List<? extends OffsetBasedSource<T>> split(
      long desiredBundleSizeBytes, PipelineOptions options) throws Exception {
    // Split the range into bundles based on the desiredBundleSizeBytes. If the desired bundle
    // size is smaller than the minBundleSize of the source then minBundleSize will be used instead.

    long desiredBundleSizeOffsetUnits = Math.max(
        Math.max(1, desiredBundleSizeBytes / getBytesPerOffset()),
        minBundleSize);

    List<OffsetBasedSource<T>> subSources = new ArrayList<>();
    for (OffsetRange range :
        new OffsetRange(startOffset, Math.min(endOffset, getMaxEndOffset(options)))
            .split(desiredBundleSizeOffsetUnits, minBundleSize)) {
      subSources.add(createSourceForSubrange(range.getFrom(), range.getTo()));
    }
    return subSources;
  }

  @Override
  public void validate() {
    checkArgument(
        this.startOffset >= 0,
        "Start offset has value %s, must be non-negative", this.startOffset);
    checkArgument(
        this.endOffset >= 0,
        "End offset has value %s, must be non-negative", this.endOffset);
    checkArgument(
        this.startOffset <= this.endOffset,
        "Start offset %s may not be larger than end offset %s",
        this.startOffset, this.endOffset);
    checkArgument(
        this.minBundleSize >= 0,
        "minBundleSize has value %s, must be non-negative",
        this.minBundleSize);
  }

  @Override
  public String toString() {
    return "[" + startOffset + ", " + endOffset + ")";
  }

  /**
   * Returns approximately how many bytes of data correspond to a single offset in this source.
   * Used for translation between this source's range and methods defined in terms of bytes, such
   * as {@link #getEstimatedSizeBytes} and {@link #split}.
   *
   * <p>Defaults to {@code 1} byte, which is the common case for, e.g., file sources.
   */
  public long getBytesPerOffset() {
    return 1L;
  }

  /**
   * Returns the actual ending offset of the current source. The value returned by this function
   * will be used to clip the end of the range {@code [startOffset, endOffset)} such that the
   * range used is {@code [startOffset, min(endOffset, maxEndOffset))}.
   *
   * <p>As an example in which {@link OffsetBasedSource} is used to implement a file source, suppose
   * that this source was constructed with an {@code endOffset} of {@link Long#MAX_VALUE} to
   * indicate that a file should be read to the end. Then this function should determine
   * the actual, exact size of the file in bytes and return it.
   */
  public abstract long getMaxEndOffset(PipelineOptions options) throws Exception;

  /**
   * Returns an {@link OffsetBasedSource} for a subrange of the current source. The
   * subrange {@code [start, end)} must be within the range {@code [startOffset, endOffset)} of
   * the current source, i.e. {@code startOffset <= start < end <= endOffset}.
   */
  public abstract OffsetBasedSource<T> createSourceForSubrange(long start, long end);

  @Override
  public void populateDisplayData(DisplayData.Builder builder) {
    super.populateDisplayData(builder);
    builder
        .addIfNotDefault(DisplayData.item("minBundleSize", minBundleSize)
            .withLabel("Minimum Bundle Size"), 1L)
        .addIfNotDefault(DisplayData.item("startOffset", startOffset)
            .withLabel("Start Read Offset"), 0L)
        .addIfNotDefault(DisplayData.item("endOffset", endOffset)
            .withLabel("End Read Offset"), Long.MAX_VALUE);
  }

  /**
   * A {@link Source.Reader} that implements code common to readers of all
   * {@link OffsetBasedSource}s.
   *
   * <p>Subclasses have to implement:
   * <ul>
   *   <li>The methods {@link #startImpl} and {@link #advanceImpl} for reading the
   *   first or subsequent records.
   *   <li>The methods {@link #getCurrent}, {@link #getCurrentOffset}, and optionally
   *   {@link #isAtSplitPoint} and {@link #getCurrentTimestamp} to access properties of
   *   the last record successfully read by {@link #startImpl} or {@link #advanceImpl}.
   * </ul>
   */
  public abstract static class OffsetBasedReader<T> extends BoundedReader<T> {
    private static final Logger LOG = LoggerFactory.getLogger(OffsetBasedReader.class);
    private OffsetBasedSource<T> source;

    /**
     * Returns true if the last call to {@link #start} or {@link #advance} returned false.
     */
    public final boolean isDone() {
      return rangeTracker.isDone();
    }

    /**
     * Returns true if there has been a call to {@link #start}.
     */
    public final boolean isStarted() {
      return rangeTracker.isStarted();
    }

    /** The {@link OffsetRangeTracker} managing the range and current position of the source. */
    private final OffsetRangeTracker rangeTracker;

    /**
     * @param source the {@link OffsetBasedSource} to be read by the current reader.
     */
    public OffsetBasedReader(OffsetBasedSource<T> source) {
      this.source = source;
      this.rangeTracker = new OffsetRangeTracker(source.getStartOffset(), source.getEndOffset());
    }

    /**
     * Returns the <i>starting</i> offset of the {@link Source.Reader#getCurrent current record},
     * which has been read by the last successful {@link Source.Reader#start} or
     * {@link Source.Reader#advance} call.
     *
     * <p>If no such call has been made yet, the return value is unspecified.
     *
     * <p>See {@link RangeTracker} for description of offset semantics.
     */
    protected abstract long getCurrentOffset() throws NoSuchElementException;

    /**
     * Returns whether the current record is at a split point (i.e., whether the current record
     * would be the first record to be read by a source with a specified start offset of
     * {@link #getCurrentOffset}).
     *
     * <p>See detailed documentation about split points in {@link RangeTracker}.
     */
    protected boolean isAtSplitPoint() throws NoSuchElementException {
      return true;
    }

    @Override
    public final boolean start() throws IOException {
      return (startImpl() && rangeTracker.tryReturnRecordAt(isAtSplitPoint(), getCurrentOffset()))
          || rangeTracker.markDone();
    }

    @Override
    public final boolean advance() throws IOException {
      return (advanceImpl() && rangeTracker.tryReturnRecordAt(isAtSplitPoint(), getCurrentOffset()))
          || rangeTracker.markDone();
    }

    /**
     * Initializes the {@link OffsetBasedSource.OffsetBasedReader} and advances to the first record,
     * returning {@code true} if there is a record available to be read. This method will be
     * invoked exactly once and may perform expensive setup operations that are needed to
     * initialize the reader.
     *
     * <p>This function is the {@code OffsetBasedReader} implementation of
     * {@link BoundedReader#start}. The key difference is that the implementor can ignore the
     * possibility that it should no longer produce the first record, either because it has exceeded
     * the original {@code endOffset} assigned to the reader, or because a concurrent call to
     * {@link #splitAtFraction} has changed the source to shrink the offset range being read.
     *
     * @see BoundedReader#start
     */
    protected abstract boolean startImpl() throws IOException;

    /**
     * Advances to the next record and returns {@code true}, or returns false if there is no next
     * record.
     *
     * <p>This function is the {@code OffsetBasedReader} implementation of
     * {@link BoundedReader#advance}. The key difference is that the implementor can ignore the
     * possibility that it should no longer produce the next record, either because it has exceeded
     * the original {@code endOffset} assigned to the reader, or because a concurrent call to
     * {@link #splitAtFraction} has changed the source to shrink the offset range being read.
     *
     * @see BoundedReader#advance
     */
    protected abstract boolean advanceImpl() throws IOException;

    @Override
    public synchronized OffsetBasedSource<T> getCurrentSource() {
      return source;
    }

    @Override
    public Double getFractionConsumed() {
      return rangeTracker.getFractionConsumed();
    }

    @Override
    public long getSplitPointsConsumed() {
      return rangeTracker.getSplitPointsProcessed();
    }

    @Override
    public long getSplitPointsRemaining() {
      if (isDone()) {
        return 0;
      } else if (!isStarted()) {
        // Note that even if the current source does not allow splitting, we don't know that
        // it's non-empty so we return UNKNOWN instead of 1.
        return BoundedReader.SPLIT_POINTS_UNKNOWN;
      } else if (!allowsDynamicSplitting()) {
        // Started (so non-empty) and unsplittable, so only the current task.
        return 1;
      } else if (getCurrentOffset() >= rangeTracker.getStopPosition() - 1) {
        // If this is true, the next element is outside the range. Note that even getCurrentOffset()
        // might be larger than the stop position when the current record is not a split point.
        return 1;
      } else {
        // Use the default.
        return super.getSplitPointsRemaining();
      }
    }

    /**
     * Whether this reader should allow dynamic splitting of the offset ranges.
     *
     * <p>True by default. Override this to return false if the reader cannot
     * support dynamic splitting correctly. If this returns false,
     * {@link OffsetBasedReader#splitAtFraction} will refuse all split requests.
     */
    public boolean allowsDynamicSplitting() {
      return true;
    }

    @Override
    public final synchronized OffsetBasedSource<T> splitAtFraction(double fraction) {
      if (!allowsDynamicSplitting()) {
        return null;
      }
      if (rangeTracker.getStopPosition() == Long.MAX_VALUE) {
        LOG.debug(
            "Refusing to split unbounded OffsetBasedReader {} at fraction {}",
            rangeTracker, fraction);
        return null;
      }
      long splitOffset = rangeTracker.getPositionForFractionConsumed(fraction);
      LOG.debug(
          "Proposing to split OffsetBasedReader {} at fraction {} (offset {})",
          rangeTracker, fraction, splitOffset);
      long start = source.getStartOffset();
      long end = source.getEndOffset();
      OffsetBasedSource<T> primary = source.createSourceForSubrange(start, splitOffset);
      OffsetBasedSource<T> residual = source.createSourceForSubrange(splitOffset, end);
      if (!rangeTracker.trySplitAtPosition(splitOffset)) {
        return null;
      }
      this.source = primary;
      return residual;
    }
  }
}
