/*
 * 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.dataflow.worker.util.common.worker;

import java.io.IOException;
import java.util.NoSuchElementException;
import java.util.Objects;
import java.util.Observable;
import org.checkerframework.checker.nullness.qual.Nullable;

/**
 * Abstract base class for native readers in the Dataflow runner.
 *
 * <p>A {@link com.google.api.services.dataflow.model.Source} is read from by getting an {@code
 * Iterator}-like value and iterating through it.
 *
 * <p>This class is intended for formats that have built-in support on the Dataflow service. <b>Do
 * not introduce new implementations:</b> for creating new input formats, use {@code
 * Bounded/UnboundedSource} instead.
 *
 * @param <T> the type of the elements read from the source
 */
public abstract class NativeReader<T> extends Observable {

  /** Returns a ReaderIterator that allows reading from this source. */
  public abstract NativeReaderIterator<T> iterator() throws IOException;

  /**
   * A stateful iterator over the data in a {@link NativeReader}.
   *
   * <p>This class is only partially thread safe. Methods {@link #start}, {@link #getCurrent},
   * {@link #advance}, and {@link #close} must be called serially. Methods {@link
   * #requestCheckpoint} and {@link #requestDynamicSplit} may be called asynchronously to those
   * methods but must not be called concurrently to themselves or each other. Methods {@link
   * #getProgress} and {@link #getRemainingParallelism} may be called asynchronously to any of the
   * methods including themselves.
   */
  public abstract static class NativeReaderIterator<T> implements AutoCloseable {

    /** A default value to return from {@link #getRemainingParallelism()}. */
    public static final double REMAINING_PARALLELISM_FROM_PROGRESS_FRACTION = Double.NaN;

    /**
     * Returns a representation of how far this iterator is through the source.
     *
     * @return the progress, or {@code null} if no progress measure can be provided (implementors
     *     are discouraged from throwing {@code UnsupportedOperationException} in this case). By
     *     default, returns {@code null}.
     */
    public @Nullable Progress getProgress() {
      return null;
    }

    /**
     * Like {@link #requestDynamicSplit}, but rather than taking the split point, makes the
     * "primary" part as small as possible, so that this iterator will produce a minimal number of
     * (ideally no) further records.
     */
    public @Nullable DynamicSplitResult requestCheckpoint() {
      return null;
    }

    /**
     * Attempts to split the input in two parts: the "primary" part and the "residual" part. The
     * current {@link NativeReaderIterator} keeps processing the primary part, while the residual
     * part will be processed elsewhere (e.g. perhaps on a different worker).
     *
     * <p>The primary and residual parts, if concatenated, must represent the same input as the
     * current input of this {@link NativeReaderIterator} before this call.
     *
     * <p>The boundary between the primary part and the residual part is specified in a
     * framework-specific way using {@link NativeReader.DynamicSplitRequest}: e.g., if the framework
     * supports the notion of positions, it might be a position at which the input is asked to split
     * itself (which is not necessarily the same position at which it <i>will</i> split itself); it
     * might be an approximate fraction of input, or something else.
     *
     * <p>{@link NativeReader.DynamicSplitResult} encodes, in a framework-specific way, the
     * information sufficient to construct a description of the resulting primary and residual
     * inputs. For example, it might, again, be a position demarcating these parts, or it might be a
     * pair of fully-specified input descriptions, or something else.
     *
     * <p>After a successful call to {@link #requestDynamicSplit}, subsequent calls should be
     * interpreted relative to the new primary.
     *
     * <p>This call should not affect the range of input represented by the {@link NativeReader}
     * that produced this {@link NativeReaderIterator}.
     *
     * @return {@code null} if the {@link NativeReader.DynamicSplitRequest} cannot be honored (in
     *     that case the input represented by this {@link NativeReaderIterator} stays the same), or
     *     a {@link NativeReader.DynamicSplitResult} describing how the input was split into a
     *     primary and residual part. By default, returns {@code null}.
     */
    public @Nullable DynamicSplitResult requestDynamicSplit(DynamicSplitRequest request) {
      return null;
    }

    /**
     * Returns an estimate of the degree of parallelism that could be achieved by {@link
     * #requestDynamicSplit} taking into account what has already been consumed. E.g., if the reader
     * has just returned the last record in the source, the remaining parallelism is 1 because it
     * can't be split up any further. If the reader just returned the 3rd record in a perfectly
     * parallelizable source with 5 records, the remaining parallelism is 3 because it could be
     * processed in parallel by this worker and two others. If the reader does not support dynamic
     * splitting, the remaining parallelism is always 1.
     *
     * <p>An exact number isn't required, mostly we want to be able to distinguish between many,
     * few, or one. Should not block.
     *
     * <p>An implementor may return {@link #REMAINING_PARALLELISM_FROM_PROGRESS_FRACTION}, in which
     * case the remaining parallelism will be interpolated from consumed parallelism, if available,
     * and progress fraction. {@link Double#POSITIVE_INFINITY} may also be returned (indicating no
     * known bound on parallelism), as may fractional estimates (in which case the sum over all
     * shards is taken).
     *
     * <p>By default, returns {@link #REMAINING_PARALLELISM_FROM_PROGRESS_FRACTION}.
     */
    public double getRemainingParallelism() {
      return REMAINING_PARALLELISM_FROM_PROGRESS_FRACTION;
    }

    /**
     * Initializes the reader and advances the reader to the first record.
     *
     * <p>This method should be called exactly once. The invocation should occur prior to calling
     * {@link #advance} or {@link #getCurrent}. This method may perform expensive operations that
     * are needed to initialize the reader.
     *
     * @return {@code true} if a record was read, {@code false} if there is no more input available.
     */
    public abstract boolean start() throws IOException;

    /**
     * Advances the reader to the next valid record.
     *
     * <p>It is an error to call this without having called {@link #start} first.
     *
     * @return {@code true} if a record was read, {@code false} if there is no more input available.
     */
    public abstract boolean advance() throws IOException;

    /**
     * Returns the value of the data item that was read by the last {@link #start} or {@link
     * #advance} call. The returned value must be effectively immutable and remain valid
     * indefinitely.
     *
     * <p>Multiple calls to this method without an intervening call to {@link #advance} should
     * return the same result.
     *
     * @throws NoSuchElementException if {@link #start} was never called, or if the last {@link
     *     #start} or {@link #advance} returned {@code false}
     */
    public abstract T getCurrent() throws NoSuchElementException;

    @Override
    public void close() throws IOException {
      // By default, do nothing.
    }

    public void abort() throws IOException {
      // By default, just close
      close();
    }

    /*
     * Called if the reader must be notified that it ought to abort. Overrides must ensure thread
     * safety as this method may not be called from the same thread, or serialized with, calls to
     * the rest of the API.
     */
    public void asyncAbort() {
      // By default, do nothing.
    }
  }

  /**
   * A representation of how far a {@code ReaderIterator} is through a {@code Reader}.
   *
   * <p>The common worker framework does not interpret instances of this interface. But a
   * tool-specific framework can make assumptions about the implementation, and so the concrete
   * Reader subclasses used by a tool-specific framework should match.
   */
  public interface Progress {}

  /**
   * A representation of a position in an iteration through a {@code Reader}.
   *
   * <p>See the comment on {@link Progress} for how instances of this interface are used by the rest
   * of the framework.
   */
  public interface Position {}

  /**
   * A framework-specific way to specify how {@link NativeReaderIterator#requestDynamicSplit} should
   * split the input into a primary and residual part.
   */
  public interface DynamicSplitRequest {}

  /**
   * A framework-specific way to specify how {@link NativeReaderIterator#requestDynamicSplit} has
   * split the input into a primary and residual part.
   */
  public interface DynamicSplitResult {}

  /**
   * A {@link NativeReader.DynamicSplitResult} that specifies the boundary between the primary and
   * residual parts of the input using a {@link Position}.
   */
  public static final class DynamicSplitResultWithPosition implements DynamicSplitResult {
    private final Position acceptedPosition;

    public DynamicSplitResultWithPosition(Position acceptedPosition) {
      this.acceptedPosition = acceptedPosition;
    }

    public Position getAcceptedPosition() {
      return acceptedPosition;
    }

    @Override
    public String toString() {
      return String.valueOf(acceptedPosition);
    }

    @Override
    public int hashCode() {
      return acceptedPosition.hashCode();
    }

    @Override
    public boolean equals(@Nullable Object obj) {
      if (this == obj) {
        return true;
      } else if (!(obj instanceof DynamicSplitResultWithPosition)) {
        return false;
      } else {
        return Objects.equals(
            acceptedPosition, ((DynamicSplitResultWithPosition) obj).acceptedPosition);
      }
    }
  }

  /**
   * Utility method to notify observers about a new element, which has been read by this Reader, and
   * its size in bytes. Normally, there is only one observer, which is a ReadOperation that
   * encapsules this Reader. Derived classes must call this method whenever they read additional
   * data, even if that element may never be returned from the corresponding source iterator.
   */
  protected void notifyElementRead(long byteSize) {
    setChanged();
    notifyObservers(byteSize);
  }

  /** Returns whether this Reader can be restarted. */
  public boolean supportsRestart() {
    return false;
  }
}
