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

import static com.google.cloud.pubsublite.internal.Preconditions.checkState;

import com.google.cloud.pubsublite.Offset;
import com.google.cloud.pubsublite.Partition;
import com.google.cloud.pubsublite.internal.BufferingPullSubscriber;
import com.google.cloud.pubsublite.internal.wire.Committer;
import com.google.cloud.pubsublite.internal.wire.SubscriberFactory;
import com.google.cloud.pubsublite.proto.Cursor;
import com.google.cloud.pubsublite.proto.SeekRequest;
import com.google.cloud.pubsublite.proto.SequencedMessage;
import io.grpc.StatusException;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.extensions.protobuf.ProtoCoder;
import org.apache.beam.sdk.io.UnboundedSource;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
import org.checkerframework.checker.nullness.qual.Nullable;

/** An UnboundedSource of Pub/Sub Lite SequencedMessages. */
class PubsubLiteUnboundedSource extends UnboundedSource<SequencedMessage, OffsetCheckpointMark> {
  private final SubscriberOptions subscriberOptions;

  PubsubLiteUnboundedSource(SubscriberOptions options) {
    this.subscriberOptions = options;
  }

  @Override
  public List<? extends UnboundedSource<SequencedMessage, OffsetCheckpointMark>> split(
      int desiredNumSplits, PipelineOptions options) {
    ArrayList<ArrayList<Partition>> partitionPartitions = new ArrayList<>(desiredNumSplits);
    for (int i = 0; i < desiredNumSplits; i++) {
      partitionPartitions.add(new ArrayList<>());
    }
    int counter = 0;
    for (Partition partition : subscriberOptions.partitions()) {
      partitionPartitions.get(counter % desiredNumSplits).add(partition);
      ++counter;
    }
    ImmutableList.Builder<PubsubLiteUnboundedSource> builder = ImmutableList.builder();
    for (List<Partition> partitionSubset : partitionPartitions) {
      if (partitionSubset.isEmpty()) {
        continue;
      }
      try {
        builder.add(
            new PubsubLiteUnboundedSource(
                subscriberOptions
                    .toBuilder()
                    .setPartitions(ImmutableSet.copyOf(partitionSubset))
                    .build()));
      } catch (StatusException e) {
        throw e.getStatus().asRuntimeException();
      }
    }
    return builder.build();
  }

  @Override
  public UnboundedReader<SequencedMessage> createReader(
      PipelineOptions options, @Nullable OffsetCheckpointMark checkpointMark) throws IOException {
    try {
      Map<Partition, SubscriberFactory> subscriberFactories =
          subscriberOptions.getSubscriberFactories();
      Map<Partition, Committer> committers = subscriberOptions.getCommitters();
      ImmutableMap.Builder<Partition, PubsubLiteUnboundedReader.SubscriberState> statesBuilder =
          ImmutableMap.builder();
      for (Partition partition : subscriberFactories.keySet()) {
        checkState(committers.containsKey(partition));
        PubsubLiteUnboundedReader.SubscriberState state =
            new PubsubLiteUnboundedReader.SubscriberState();
        state.committer = committers.get(partition);
        if (checkpointMark != null && checkpointMark.partitionOffsetMap.containsKey(partition)) {
          Offset checkpointed = checkpointMark.partitionOffsetMap.get(partition);
          state.lastDelivered = Optional.of(checkpointed);
          state.subscriber =
              new TranslatingPullSubscriber(
                  new BufferingPullSubscriber(
                      subscriberFactories.get(partition),
                      subscriberOptions.flowControlSettings(),
                      SeekRequest.newBuilder()
                          .setCursor(Cursor.newBuilder().setOffset(checkpointed.value()))
                          .build()));
        } else {
          state.subscriber =
              new TranslatingPullSubscriber(
                  new BufferingPullSubscriber(
                      subscriberFactories.get(partition), subscriberOptions.flowControlSettings()));
        }
        statesBuilder.put(partition, state);
      }
      return new PubsubLiteUnboundedReader(this, statesBuilder.build());
    } catch (StatusException e) {
      throw new IOException(e);
    }
  }

  @Override
  public Coder<OffsetCheckpointMark> getCheckpointMarkCoder() {
    return OffsetCheckpointMark.getCoder();
  }

  @Override
  public Coder<SequencedMessage> getOutputCoder() {
    return ProtoCoder.of(SequencedMessage.class);
  }
}

@@ -923,16 +923,50 @@
 ckState;
+%0Aimport static java.lang.Math.min;
 %0A%0Aimport
@@ -2628,16 +2628,29 @@
 yList%3C%3E(
+%0A        min(
 desiredN
@@ -2657,16 +2657,56 @@
 umSplits
+, subscriberOptions.partitions().size())
 );%0A    f

Fix nit