/*
 * 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.core.construction.graph;

import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;

import com.google.auto.value.AutoValue;
import java.util.ArrayList;
import java.util.Collection;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.function.Predicate;
import java.util.stream.Collectors;
import javax.annotation.Nullable;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.model.pipeline.v1.RunnerApi.Components;
import org.apache.beam.model.pipeline.v1.RunnerApi.FunctionSpec;
import org.apache.beam.model.pipeline.v1.RunnerApi.PCollection;
import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
import org.apache.beam.runners.core.construction.PTransformTranslation;
import org.apache.beam.runners.core.construction.SyntheticComponents;
import org.apache.beam.runners.core.construction.graph.PipelineNode.PCollectionNode;
import org.apache.beam.runners.core.construction.graph.PipelineNode.PTransformNode;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.HashMultimap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Multimap;

/**
 * Utilities to insert synthetic {@link PCollectionNode PCollections} for {@link PCollection
 * PCollections} which are produced by multiple independently executable stages.
 */
class OutputDeduplicator {

  /**
   * Ensure that no {@link PCollection} output by any of the {@code stages} or {@code
   * unfusedTransforms} is produced by more than one of those stages or transforms.
   *
   * <p>For each {@link PCollection} output by multiple stages and/or transforms, each producer is
   * rewritten to produce a partial {@link PCollection}, which are then flattened together via an
   * introduced Flatten node which produces the original output.
   */
  static DeduplicationResult ensureSingleProducer(
      QueryablePipeline pipeline,
      Collection<ExecutableStage> stages,
      Collection<PTransformNode> unfusedTransforms) {
    RunnerApi.Components.Builder unzippedComponents = pipeline.getComponents().toBuilder();

    Multimap<PCollectionNode, StageOrTransform> pcollectionProducers =
        getProducers(pipeline, stages, unfusedTransforms);
    Multimap<StageOrTransform, PCollectionNode> requiresNewOutput = HashMultimap.create();
    // Create a synthetic PCollection for each of these nodes. The transforms in the runner
    // portion of the graph that creates them should be replaced in the result components. The
    // ExecutableStage must also be rewritten to have updated outputs and transforms.
    for (Map.Entry<PCollectionNode, Collection<StageOrTransform>> collectionProducer :
        pcollectionProducers.asMap().entrySet()) {
      if (collectionProducer.getValue().size() > 1) {
        for (StageOrTransform producer : collectionProducer.getValue()) {
          requiresNewOutput.put(producer, collectionProducer.getKey());
        }
      }
    }

    Map<ExecutableStage, ExecutableStage> updatedStages = new LinkedHashMap<>();
    Map<String, PTransformNode> updatedTransforms = new LinkedHashMap<>();
    Multimap<String, PCollectionNode> originalToPartial = HashMultimap.create();
    for (Map.Entry<StageOrTransform, Collection<PCollectionNode>> deduplicationTargets :
        requiresNewOutput.asMap().entrySet()) {
      if (deduplicationTargets.getKey().getStage() != null) {
        StageDeduplication deduplication =
            deduplicatePCollections(
                deduplicationTargets.getKey().getStage(),
                deduplicationTargets.getValue(),
                unzippedComponents::containsPcollections);
        for (Entry<String, PCollectionNode> originalToPartialReplacement :
            deduplication.getOriginalToPartialPCollections().entrySet()) {
          originalToPartial.put(
              originalToPartialReplacement.getKey(), originalToPartialReplacement.getValue());
          unzippedComponents.putPcollections(
              originalToPartialReplacement.getValue().getId(),
              originalToPartialReplacement.getValue().getPCollection());
        }
        updatedStages.put(
            deduplicationTargets.getKey().getStage(), deduplication.getUpdatedStage());
      } else if (deduplicationTargets.getKey().getTransform() != null) {
        PTransformDeduplication deduplication =
            deduplicatePCollections(
                deduplicationTargets.getKey().getTransform(),
                deduplicationTargets.getValue(),
                unzippedComponents::containsPcollections);
        for (Entry<String, PCollectionNode> originalToPartialReplacement :
            deduplication.getOriginalToPartialPCollections().entrySet()) {
          originalToPartial.put(
              originalToPartialReplacement.getKey(), originalToPartialReplacement.getValue());
          unzippedComponents.putPcollections(
              originalToPartialReplacement.getValue().getId(),
              originalToPartialReplacement.getValue().getPCollection());
        }
        updatedTransforms.put(
            deduplicationTargets.getKey().getTransform().getId(),
            deduplication.getUpdatedTransform());
      } else {
        throw new IllegalStateException(
            String.format(
                "%s with no %s or %s",
                StageOrTransform.class.getSimpleName(),
                ExecutableStage.class.getSimpleName(),
                PTransformNode.class.getSimpleName()));
      }
    }

    Set<PTransformNode> introducedFlattens = new LinkedHashSet<>();
    for (Map.Entry<String, Collection<PCollectionNode>> partialFlattenTargets :
        originalToPartial.asMap().entrySet()) {
      String flattenId =
          SyntheticComponents.uniqueId("unzipped_flatten", unzippedComponents::containsTransforms);
      PTransform flattenPartialPCollections =
          createFlattenOfPartials(
              flattenId, partialFlattenTargets.getKey(), partialFlattenTargets.getValue());
      unzippedComponents.putTransforms(flattenId, flattenPartialPCollections);
      introducedFlattens.add(PipelineNode.pTransform(flattenId, flattenPartialPCollections));
    }

    Components components = unzippedComponents.build();
    return DeduplicationResult.of(components, introducedFlattens, updatedStages, updatedTransforms);
  }

  @AutoValue
  abstract static class DeduplicationResult {
    private static DeduplicationResult of(
        RunnerApi.Components components,
        Set<PTransformNode> introducedTransforms,
        Map<ExecutableStage, ExecutableStage> stages,
        Map<String, PTransformNode> unfused) {
      return new AutoValue_OutputDeduplicator_DeduplicationResult(
          components, introducedTransforms, stages, unfused);
    }

    abstract RunnerApi.Components getDeduplicatedComponents();

    abstract Set<PTransformNode> getIntroducedTransforms();

    abstract Map<ExecutableStage, ExecutableStage> getDeduplicatedStages();

    abstract Map<String, PTransformNode> getDeduplicatedTransforms();
  }

  private static PTransform createFlattenOfPartials(
      String transformId, String outputId, Collection<PCollectionNode> generatedInputs) {
    PTransform.Builder newFlattenBuilder = PTransform.newBuilder();
    int i = 0;
    for (PCollectionNode generatedInput : generatedInputs) {
      String localInputId = String.format("input_%s", i);
      i++;
      newFlattenBuilder.putInputs(localInputId, generatedInput.getId());
    }
    // Flatten all of the new partial nodes together.
    return newFlattenBuilder
        // Use transform ID as unique name.
        .setUniqueName(transformId)
        .putOutputs("output", outputId)
        .setSpec(FunctionSpec.newBuilder().setUrn(PTransformTranslation.FLATTEN_TRANSFORM_URN))
        .build();
  }

  /**
   * Returns the map from each {@link PCollectionNode} produced by any of the {@link ExecutableStage
   * stages} or {@link PTransformNode transforms} to all of the {@link ExecutableStage stages} or
   * {@link PTransformNode transforms} that produce it.
   */
  private static Multimap<PCollectionNode, StageOrTransform> getProducers(
      QueryablePipeline pipeline,
      Iterable<ExecutableStage> stages,
      Iterable<PTransformNode> unfusedTransforms) {
    Multimap<PCollectionNode, StageOrTransform> pcollectionProducers = HashMultimap.create();
    for (ExecutableStage stage : stages) {
      for (PCollectionNode output : stage.getOutputPCollections()) {
        pcollectionProducers.put(output, StageOrTransform.stage(stage));
      }
    }
    for (PTransformNode unfused : unfusedTransforms) {
      for (PCollectionNode output : pipeline.getOutputPCollections(unfused)) {
        pcollectionProducers.put(output, StageOrTransform.transform(unfused));
      }
    }
    return pcollectionProducers;
  }

  private static PTransformDeduplication deduplicatePCollections(
      PTransformNode transform,
      Collection<PCollectionNode> duplicates,
      Predicate<String> existingPCollectionIds) {
    Map<String, PCollectionNode> unzippedOutputs =
        createPartialPCollections(duplicates, existingPCollectionIds);
    PTransform pTransform = updateOutputs(transform.getTransform(), unzippedOutputs);
    return PTransformDeduplication.of(
        PipelineNode.pTransform(transform.getId(), pTransform), unzippedOutputs);
  }

  @AutoValue
  abstract static class PTransformDeduplication {
    public static PTransformDeduplication of(
        PTransformNode updatedTransform, Map<String, PCollectionNode> originalToPartial) {
      return new AutoValue_OutputDeduplicator_PTransformDeduplication(
          updatedTransform, originalToPartial);
    }

    abstract PTransformNode getUpdatedTransform();

    abstract Map<String, PCollectionNode> getOriginalToPartialPCollections();
  }

  private static StageDeduplication deduplicatePCollections(
      ExecutableStage stage,
      Collection<PCollectionNode> duplicates,
      Predicate<String> existingPCollectionIds) {
    Map<String, PCollectionNode> unzippedOutputs =
        createPartialPCollections(duplicates, existingPCollectionIds);
    ExecutableStage updatedStage = deduplicateStageOutput(stage, unzippedOutputs);
    return StageDeduplication.of(updatedStage, unzippedOutputs);
  }

  @AutoValue
  abstract static class StageDeduplication {
    public static StageDeduplication of(
        ExecutableStage updatedStage, Map<String, PCollectionNode> originalToPartial) {
      return new AutoValue_OutputDeduplicator_StageDeduplication(updatedStage, originalToPartial);
    }

    abstract ExecutableStage getUpdatedStage();

    abstract Map<String, PCollectionNode> getOriginalToPartialPCollections();
  }

  /**
   * Returns a {@link Map} from the ID of a {@link PCollectionNode PCollection} to a {@link
   * PCollectionNode} that contains part of that {@link PCollectionNode PCollection}.
   */
  private static Map<String, PCollectionNode> createPartialPCollections(
      Collection<PCollectionNode> duplicates, Predicate<String> existingPCollectionIds) {
    Map<String, PCollectionNode> unzippedOutputs = new LinkedHashMap<>();
    Predicate<String> existingOrNewIds =
        existingPCollectionIds.or(
            id ->
                unzippedOutputs.values().stream().map(PCollectionNode::getId).anyMatch(id::equals));
    for (PCollectionNode duplicateOutput : duplicates) {
      String id = SyntheticComponents.uniqueId(duplicateOutput.getId(), existingOrNewIds);
      PCollection partial = duplicateOutput.getPCollection().toBuilder().setUniqueName(id).build();
      // Check to make sure there is only one duplicated output with the same id - which ensures we
      // only introduce one 'partial output' per producer of that output.
      PCollectionNode alreadyDeduplicated =
          unzippedOutputs.put(duplicateOutput.getId(), PipelineNode.pCollection(id, partial));
      checkArgument(alreadyDeduplicated == null, "a duplicate should only appear once per stage");
    }
    return unzippedOutputs;
  }

  /**
   * Returns an {@link ExecutableStage} where all of the {@link PCollectionNode PCollections}
   * matching the original are replaced with the introduced partial {@link PCollection} in all
   * references made within the {@link ExecutableStage}.
   */
  private static ExecutableStage deduplicateStageOutput(
      ExecutableStage stage, Map<String, PCollectionNode> originalToPartial) {
    Collection<PTransformNode> updatedTransforms = new ArrayList<>();
    for (PTransformNode transform : stage.getTransforms()) {
      PTransform updatedTransform = updateOutputs(transform.getTransform(), originalToPartial);
      updatedTransforms.add(PipelineNode.pTransform(transform.getId(), updatedTransform));
    }
    Collection<PCollectionNode> updatedOutputs = new ArrayList<>();
    for (PCollectionNode output : stage.getOutputPCollections()) {
      updatedOutputs.add(originalToPartial.getOrDefault(output.getId(), output));
    }
    RunnerApi.Components updatedStageComponents =
        stage
            .getComponents()
            .toBuilder()
            .clearTransforms()
            .putAllTransforms(
                updatedTransforms.stream()
                    .collect(Collectors.toMap(PTransformNode::getId, PTransformNode::getTransform)))
            .putAllPcollections(
                originalToPartial.values().stream()
                    .collect(
                        Collectors.toMap(PCollectionNode::getId, PCollectionNode::getPCollection)))
            .build();
    return ImmutableExecutableStage.of(
        updatedStageComponents,
        stage.getEnvironment(),
        stage.getInputPCollection(),
        stage.getSideInputs(),
        stage.getUserStates(),
        stage.getTimers(),
        updatedTransforms,
        updatedOutputs);
  }

  /**
   * Returns a {@link PTransform} like the input {@link PTransform}, but with each output to {@code
   * originalPCollection} replaced with an output (with the same local name) to {@code
   * newPCollection}.
   */
  private static PTransform updateOutputs(
      PTransform transform, Map<String, PCollectionNode> originalToPartial) {
    PTransform.Builder updatedTransformBuilder = transform.toBuilder();
    for (Map.Entry<String, String> output : transform.getOutputsMap().entrySet()) {
      if (originalToPartial.containsKey(output.getValue())) {
        updatedTransformBuilder.putOutputs(
            output.getKey(), originalToPartial.get(output.getValue()).getId());
      }
    }
    return updatedTransformBuilder.build();
  }

  @AutoValue
  abstract static class StageOrTransform {
    public static StageOrTransform stage(ExecutableStage stage) {
      return new AutoValue_OutputDeduplicator_StageOrTransform(stage, null);
    }

    public static StageOrTransform transform(PTransformNode transform) {
      return new AutoValue_OutputDeduplicator_StageOrTransform(null, transform);
    }

    @Nullable
    abstract ExecutableStage getStage();

    @Nullable
    abstract PTransformNode getTransform();
  }
}