/*
 * Copyright 2020 Google LLC
 *
 * Licensed 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 com.google.cloud.pubsublite.beam;

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.SequencedMessage;
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.common.base.Ticker;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import io.grpc.StatusException;
import java.io.IOException;
import java.util.List;
import java.util.Optional;
import javax.annotation.Nullable;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.io.UnboundedSource;
import org.apache.beam.sdk.options.PipelineOptions;

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) {
    ImmutableList.Builder<PubsubLiteUnboundedSource> builder = ImmutableList.builder();
    for (List<Partition> partitionSubset :
        Iterables.partition(subscriberOptions.partitions(), desiredNumSplits)) {
      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 {
      ImmutableMap<Partition, SubscriberFactory> subscriberFactories =
          subscriberOptions.getSubscriberFactories();
      ImmutableMap<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 BufferingPullSubscriber(
                  subscriberFactories.get(partition),
                  subscriberOptions.flowControlSettings(),
                  SeekRequest.newBuilder()
                      .setCursor(Cursor.newBuilder().setOffset(checkpointed.value()))
                      .build());
        } else {
          state.subscriber =
              new BufferingPullSubscriber(
                  subscriberFactories.get(partition), subscriberOptions.flowControlSettings());
        }
        statesBuilder.put(partition, state);
      }
      return new PubsubLiteUnboundedReader(
          this,
          statesBuilder.build(),
          TopicBacklogReader.create(subscriberOptions.topicBacklogReaderSettings()),
          Ticker.systemTicker());
    } catch (StatusException e) {
      throw new IOException(e);
    }
  }

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

  @Override
  public Coder<SequencedMessage> getOutputCoder() {
    return new SequencedMessageCoder();
  }
}

@@ -1337,52 +1337,8 @@
 et;%0A
-import com.google.common.collect.Iterables;%0A
 impo
@@ -1385,32 +1385,60 @@
 io.IOException;%0A
+import java.util.ArrayList;%0A
 import java.util
@@ -2049,24 +2049,405 @@
  options) %7B%0A
+    ArrayList%3CArrayList%3CPartition%3E%3E partitionPartitions = new ArrayList%3C%3E(desiredNumSplits);%0A    for (int i = 0; i %3C desiredNumSplits; i++) %7B%0A      partitionPartitions.add(new ArrayList%3C%3E());%0A    %7D%0A    int counter = 0;%0A    for (Partition partition : subscriberOptions.partitions()) %7B%0A      partitionPartitions.get(counter %25 desiredNumSplits).add(partition);%0A      ++counter;%0A    %7D%0A
     Immutabl
@@ -2568,86 +2568,28 @@
 et :
-%0A        Iterables.partition(subscriberOptions.partitions(), desiredNumSplits)
+ partitionPartitions
 ) %7B%0A

Fix PubsubLiteUnboundedSource to create n partitions not partitions of n size (#313)