/*
 * 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;

import java.nio.ByteBuffer;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import org.apache.beam.runners.core.construction.PTransformReplacements;
import org.apache.beam.runners.core.construction.ReplacementOutputs;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.runners.AppliedPTransform;
import org.apache.beam.sdk.runners.PTransformOverrideFactory;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.GroupByKey;
import org.apache.beam.sdk.transforms.GroupIntoBatches;
import org.apache.beam.sdk.transforms.MapElements;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.SimpleFunction;
import org.apache.beam.sdk.util.ShardedKey;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterators;

@SuppressWarnings({
  "rawtypes" // TODO(https://issues.apache.org/jira/browse/BEAM-10556)
})
public class GroupIntoBatchesOverride {

  static class BatchGroupIntoBatchesOverrideFactory<K, V>
      implements PTransformOverrideFactory<
          PCollection<KV<K, V>>, PCollection<KV<K, Iterable<V>>>, GroupIntoBatches<K, V>> {

    @Override
    public PTransformReplacement<PCollection<KV<K, V>>, PCollection<KV<K, Iterable<V>>>>
        getReplacementTransform(
            AppliedPTransform<
                    PCollection<KV<K, V>>, PCollection<KV<K, Iterable<V>>>, GroupIntoBatches<K, V>>
                transform) {
      return PTransformReplacement.of(
          PTransformReplacements.getSingletonMainInput(transform),
          new BatchGroupIntoBatches<>(transform.getTransform().getBatchSize()));
    }

    @Override
    public Map<PCollection<?>, ReplacementOutput> mapOutputs(
        Map<TupleTag<?>, PCollection<?>> outputs, PCollection<KV<K, Iterable<V>>> newOutput) {
      return ReplacementOutputs.singleton(outputs, newOutput);
    }
  }

  /** Specialized implementation of {@link GroupIntoBatches} for bounded Dataflow pipelines. */
  static class BatchGroupIntoBatches<K, V>
      extends PTransform<PCollection<KV<K, V>>, PCollection<KV<K, Iterable<V>>>> {

    private final long batchSize;

    private BatchGroupIntoBatches(long batchSize) {
      this.batchSize = batchSize;
    }

    @Override
    public PCollection<KV<K, Iterable<V>>> expand(PCollection<KV<K, V>> input) {
      return input
          .apply("GroupAll", GroupByKey.create())
          .apply(
              "SplitIntoBatches",
              ParDo.of(
                  new DoFn<KV<K, Iterable<V>>, KV<K, Iterable<V>>>() {
                    @ProcessElement
                    public void process(ProcessContext c) {
                      // Iterators.partition lazily creates the partitions as they are accessed
                      // allowing it to partition very large iterators.
                      Iterator<List<V>> iterator =
                          Iterators.partition(c.element().getValue().iterator(), (int) batchSize);

                      // Note that GroupIntoBatches only outputs when the batch is non-empty.
                      while (iterator.hasNext()) {
                        c.output(KV.of(c.element().getKey(), iterator.next()));
                      }
                    }
                  }));
    }
  }

  static class BatchGroupIntoBatchesWithShardedKeyOverrideFactory<K, V>
      implements PTransformOverrideFactory<
          PCollection<KV<K, V>>,
          PCollection<KV<ShardedKey<K>, Iterable<V>>>,
          GroupIntoBatches<K, V>.WithShardedKey> {

    @Override
    public PTransformReplacement<PCollection<KV<K, V>>, PCollection<KV<ShardedKey<K>, Iterable<V>>>>
        getReplacementTransform(
            AppliedPTransform<
                    PCollection<KV<K, V>>,
                    PCollection<KV<ShardedKey<K>, Iterable<V>>>,
                    GroupIntoBatches<K, V>.WithShardedKey>
                transform) {
      return PTransformReplacement.of(
          PTransformReplacements.getSingletonMainInput(transform),
          new BatchGroupIntoBatchesWithShardedKey<>(transform.getTransform().getBatchSize()));
    }

    @Override
    public Map<PCollection<?>, ReplacementOutput> mapOutputs(
        Map<TupleTag<?>, PCollection<?>> outputs,
        PCollection<KV<ShardedKey<K>, Iterable<V>>> newOutput) {
      return ReplacementOutputs.singleton(outputs, newOutput);
    }
  }

  /**
   * Specialized implementation of {@link GroupIntoBatches.WithShardedKey} for bounded Dataflow
   * pipelines.
   */
  static class BatchGroupIntoBatchesWithShardedKey<K, V>
      extends PTransform<PCollection<KV<K, V>>, PCollection<KV<ShardedKey<K>, Iterable<V>>>> {

    private final long batchSize;

    private BatchGroupIntoBatchesWithShardedKey(long batchSize) {
      this.batchSize = batchSize;
    }

    @Override
    public PCollection<KV<ShardedKey<K>, Iterable<V>>> expand(PCollection<KV<K, V>> input) {
      return shardKeys(input).apply(new BatchGroupIntoBatches<>(batchSize));
    }
  }

  static class StreamingGroupIntoBatchesWithShardedKeyOverrideFactory<K, V>
      implements PTransformOverrideFactory<
          PCollection<KV<K, V>>,
          PCollection<KV<ShardedKey<K>, Iterable<V>>>,
          GroupIntoBatches<K, V>.WithShardedKey> {

    private final DataflowRunner runner;

    StreamingGroupIntoBatchesWithShardedKeyOverrideFactory(DataflowRunner runner) {
      this.runner = runner;
    }

    @Override
    public PTransformReplacement<PCollection<KV<K, V>>, PCollection<KV<ShardedKey<K>, Iterable<V>>>>
        getReplacementTransform(
            AppliedPTransform<
                    PCollection<KV<K, V>>,
                    PCollection<KV<ShardedKey<K>, Iterable<V>>>,
                    GroupIntoBatches<K, V>.WithShardedKey>
                transform) {
      return PTransformReplacement.of(
          PTransformReplacements.getSingletonMainInput(transform),
          new StreamingGroupIntoBatchesWithShardedKey<>(
              runner,
              transform.getTransform(),
              PTransformReplacements.getSingletonMainOutput(transform)));
    }

    @Override
    public Map<PCollection<?>, ReplacementOutput> mapOutputs(
        Map<TupleTag<?>, PCollection<?>> outputs,
        PCollection<KV<ShardedKey<K>, Iterable<V>>> newOutput) {
      return ReplacementOutputs.singleton(outputs, newOutput);
    }
  }

  /**
   * Specialized implementation of {@link GroupIntoBatches.WithShardedKey} for unbounded Dataflow
   * pipelines. The override does the same thing as the original transform but additionally records
   * the output in order to append required step properties during the graph translation.
   */
  static class StreamingGroupIntoBatchesWithShardedKey<K, V>
      extends PTransform<PCollection<KV<K, V>>, PCollection<KV<ShardedKey<K>, Iterable<V>>>> {

    private final transient DataflowRunner runner;
    private final GroupIntoBatches<K, V>.WithShardedKey originalTransform;
    private final transient PCollection<KV<ShardedKey<K>, Iterable<V>>> originalOutput;

    public StreamingGroupIntoBatchesWithShardedKey(
        DataflowRunner runner,
        GroupIntoBatches<K, V>.WithShardedKey original,
        PCollection<KV<ShardedKey<K>, Iterable<V>>> output) {
      this.runner = runner;
      this.originalTransform = original;
      this.originalOutput = output;
    }

    @Override
    public PCollection<KV<ShardedKey<K>, Iterable<V>>> expand(PCollection<KV<K, V>> input) {
      // Record the output PCollection of the original transform since the new output will be
      // replaced by the original one when the replacement transform is wired to other nodes in the
      // graph, although the old and the new outputs are effectively the same.
      runner.maybeRecordPCollectionWithAutoSharding(originalOutput);
      return input.apply(originalTransform);
    }
  }

  private static final UUID workerUuid = UUID.randomUUID();

  private static <K, V> PCollection<KV<ShardedKey<K>, V>> shardKeys(PCollection<KV<K, V>> input) {
    KvCoder<K, V> inputCoder = (KvCoder<K, V>) input.getCoder();
    org.apache.beam.sdk.coders.Coder<K> keyCoder =
        (org.apache.beam.sdk.coders.Coder<K>) inputCoder.getCoderArguments().get(0);
    org.apache.beam.sdk.coders.Coder<V> valueCoder =
        (org.apache.beam.sdk.coders.Coder<V>) inputCoder.getCoderArguments().get(1);
    return input
        .apply(
            "Shard Keys",
            MapElements.via(
                new SimpleFunction<KV<K, V>, KV<ShardedKey<K>, V>>() {
                  @Override
                  public KV<ShardedKey<K>, V> apply(KV<K, V> input) {
                    long tid = Thread.currentThread().getId();
                    ByteBuffer buffer = ByteBuffer.allocate(3 * Long.BYTES);
                    buffer.putLong(workerUuid.getMostSignificantBits());
                    buffer.putLong(workerUuid.getLeastSignificantBits());
                    buffer.putLong(tid);
                    return KV.of(ShardedKey.of(input.getKey(), buffer.array()), input.getValue());
                  }
                }))
        .setCoder(KvCoder.of(ShardedKey.Coder.of(keyCoder), valueCoder));
  }
}
