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

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

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.Pipeline.PipelineVisitor;
import org.apache.beam.sdk.Pipeline.PipelineVisitor.CompositeBehavior;
import org.apache.beam.sdk.annotations.Internal;
import org.apache.beam.sdk.runners.PTransformOverrideFactory.ReplacementOutput;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.resourcehints.ResourceHints;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PInput;
import org.apache.beam.sdk.values.POutput;
import org.apache.beam.sdk.values.PValue;
import org.apache.beam.sdk.values.PValues;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.MoreObjects;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * <b><i>For internal use only; no backwards-compatibility guarantees.</i></b>
 *
 * <p>Captures information about a collection of transformations and their associated {@link
 * PCollection PCollections}.
 */
@Internal
@SuppressWarnings({
  "nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-10402)
})
public class TransformHierarchy {
  private static final Logger LOG = LoggerFactory.getLogger(TransformHierarchy.class);

  private final Node root;
  private final Map<Node, PInput> unexpandedInputs;
  private final Map<PCollection<?>, Node> producers;

  // A map of PCollection<?> to the PInput the producing PTransform is applied to
  private final Map<PCollection<?>, PInput> producerInput;

  // Maintain a stack based on the enclosing nodes
  private Node current;

  public TransformHierarchy(ResourceHints resourceHints) {
    producers = new HashMap<>();
    producerInput = new HashMap<>();
    unexpandedInputs = new HashMap<>();
    root = new Node(resourceHints);
    current = root;
  }

  /**
   * Adds the named {@link PTransform} consuming the provided {@link PInput} as a node in this
   * {@link TransformHierarchy} as a child of the current node, and sets it to be the current node.
   *
   * <p>This call should be finished by expanding and recursively calling {@link #pushNode(String,
   * PInput, PTransform)}, calling {@link #finishSpecifyingInput()}, setting the output with {@link
   * #setOutput(POutput)}, and ending with a call to {@link #popNode()}.
   *
   * @return the added node
   */
  public Node pushNode(String name, PInput input, PTransform<?, ?> transform) {
    checkNotNull(
        transform, "A %s must be provided for all Nodes", PTransform.class.getSimpleName());
    checkNotNull(
        name, "A name must be provided for all %s Nodes", PTransform.class.getSimpleName());
    checkNotNull(
        input, "An input must be provided for all %s Nodes", PTransform.class.getSimpleName());
    Node node = new Node(current, transform, name, input);
    unexpandedInputs.put(node, input);
    current.addComposite(node);
    current = node;
    return current;
  }

  public Node replaceNode(Node existing, PInput input, PTransform<?, ?> transform) {
    checkNotNull(existing);
    checkNotNull(input);
    checkNotNull(transform);
    checkState(
        unexpandedInputs.isEmpty(),
        "Replacing a node when the graph has an unexpanded input. This is an SDK bug.");
    Node replacement =
        new Node(existing.getEnclosingNode(), transform, existing.getFullName(), input);
    for (PCollection<?> output : existing.getOutputs().values()) {
      Node producer = producers.get(output);
      boolean producedInExisting = false;
      do {
        if (producer.equals(existing)) {
          producedInExisting = true;
        } else {
          producer = producer.getEnclosingNode();
        }
      } while (!producedInExisting && !producer.isRootNode());
      if (producedInExisting) {
        producers.remove(output);
        LOG.debug(
            "Removed producer for value {} as it is part of a replaced composite {}",
            output,
            existing.getFullName());
      } else {
        LOG.debug("Value {} not produced in existing node {}", output, existing.getFullName());
      }
    }
    existing.getEnclosingNode().replaceChild(existing, replacement);
    unexpandedInputs.remove(existing);
    unexpandedInputs.put(replacement, input);
    current = replacement;
    return replacement;
  }

  /**
   * Finish specifying all of the input {@link PValue PValues} of the current {@link Node}. Ensures
   * that all of the inputs to the current node have been fully specified, and have been produced by
   * a node in this graph.
   */
  public void finishSpecifyingInput() {
    // Inputs must be completely specified before they are consumed by a transform.
    for (PCollection<?> inputValue : current.getInputs().values()) {
      PInput input = producerInput.remove(inputValue);
      Node producerNode = maybeGetProducer(inputValue);
      if (producerNode != null) {
        inputValue.finishSpecifying(input, producerNode.getTransform());
      }
    }
  }

  /**
   * Set the output of the current {@link Node}. If the output is new (setOutput has not previously
   * been called with it as the parameter), the current node is set as the producer of that {@link
   * POutput}.
   *
   * <p>Also validates the output - specifically, a Primitive {@link PTransform} produces all of its
   * outputs, and a Composite {@link PTransform} produces none of its outputs. Verifies that the
   * expanded output does not contain {@link PValue PValues} produced by both this node and other
   * nodes.
   */
  public void setOutput(POutput output) {
    for (PCollection<?> value : PValues.fullyExpand(output.expand()).values()) {
      if (!producers.containsKey(value)) {
        producers.put(value, current);
        value.finishSpecifyingOutput(
            current.getFullName(), unexpandedInputs.get(current), current.transform);
      }
      producerInput.put(value, unexpandedInputs.get(current));
    }
    output.finishSpecifyingOutput(
        current.getFullName(), unexpandedInputs.get(current), current.transform);
    current.setOutput(output);
  }

  /**
   * Recursively replace the outputs of the current {@link Node} with the original outputs of the
   * node it is replacing. No value that is a key in {@code originalToReplacement} may be present
   * within the {@link TransformHierarchy} after this method completes.
   */
  public void replaceOutputs(Map<PCollection<?>, ReplacementOutput> originalToReplacement) {
    current.replaceOutputs(originalToReplacement);
  }

  /**
   * Pops the current node off the top of the stack, finishing it. Outputs of the node are finished
   * once they are consumed as input.
   */
  public void popNode() {
    current.finishSpecifying();
    unexpandedInputs.remove(current);
    current = current.getEnclosingNode();
    checkState(current != null, "Can't pop the root node of a TransformHierarchy");
  }

  Node maybeGetProducer(PCollection<?> produced) {
    return producers.get(produced);
  }

  Node getProducer(PCollection<?> produced) {
    return checkNotNull(maybeGetProducer(produced), "No producer found for %s", produced);
  }

  public Set<PValue> visit(PipelineVisitor visitor) {
    finishSpecifying();
    Set<PValue> visitedValues = new HashSet<>();
    root.visit(visitor, visitedValues, new HashSet<>(), new HashSet<>());
    return visitedValues;
  }

  /**
   * Finish specifying any remaining nodes within the {@link TransformHierarchy}. These are {@link
   * PValue PValues} that are produced as output of some {@link PTransform} but are never consumed
   * as input. These values must still be finished specifying.
   */
  private void finishSpecifying() {
    for (Entry<PCollection<?>, PInput> producerInputEntry : producerInput.entrySet()) {
      PCollection<?> pCollection = producerInputEntry.getKey();
      pCollection.finishSpecifying(
          producerInputEntry.getValue(), getProducer(pCollection).getTransform());
    }
    producerInput.clear();
  }

  public Node getCurrent() {
    return current;
  }

  /**
   * Provides internal tracking of transform relationships with helper methods for initialization
   * and ordered visitation.
   */
  public class Node {
    // null for the root node, otherwise the enclosing node
    private final @Nullable Node enclosingNode;

    // The PTransform for this node, which may be a composite PTransform.
    // The root of a TransformHierarchy is represented as a Node
    // with a null transform field.
    private final @Nullable PTransform<?, ?> transform;

    private final String fullName;

    // Nodes for sub-transforms of a composite transform.
    private final List<Node> parts = new ArrayList<>();

    // Input to the transform, in expanded form.
    private final Map<TupleTag<?>, PCollection<?>> inputs;

    // TODO: track which outputs need to be exported to parent.
    // Output of the transform, in expanded form. Null if not yet set.
    private @Nullable Map<TupleTag<?>, PCollection<?>> outputs;

    private final ResourceHints resourceHints;

    @VisibleForTesting boolean finishedSpecifying = false;

    /**
     * Creates the root-level node. The root level node has a null enclosing node, a null transform,
     * an empty map of inputs, an empty map of outputs, and a name equal to the empty string.
     */
    private Node(ResourceHints resourceHints) {
      this.enclosingNode = null;
      this.transform = null;
      this.fullName = "";
      this.inputs = Collections.emptyMap();
      this.outputs = Collections.emptyMap();
      this.resourceHints = resourceHints;
    }

    /**
     * Creates a new Node with the given parent and transform.
     *
     * @param enclosingNode the composite node containing this node
     * @param transform the PTransform tracked by this node
     * @param fullName the fully qualified name of the transform
     * @param input the unexpanded input to the transform
     */
    private Node(Node enclosingNode, PTransform<?, ?> transform, String fullName, PInput input) {
      this.enclosingNode = enclosingNode;
      this.transform = transform;
      this.fullName = fullName;
      ImmutableMap.Builder<TupleTag<?>, PCollection<?>> inputs = ImmutableMap.builder();
      inputs.putAll(PValues.expandInput(input));
      inputs.putAll(PValues.fullyExpand(transform.getAdditionalInputs()));
      this.inputs = inputs.build();
      this.resourceHints = transform.getResourceHints().mergeWithOuter(enclosingNode.resourceHints);
    }

    /**
     * Creates a new {@link Node} with the given parent and transform, where inputs and outputs are
     * already known.
     *
     * <p>EnclosingNode and transform may both be null for a root-level node, which holds all other
     * nodes.
     *
     * @param enclosingNode the composite node containing this node
     * @param transform the PTransform tracked by this node
     * @param fullName the fully qualified name of the transform
     * @param inputs the expanded inputs to the transform
     * @param outputs the expanded outputs of the transform
     */
    private Node(
        @Nullable Node enclosingNode,
        @Nullable PTransform<?, ?> transform,
        String fullName,
        @Nullable Map<TupleTag<?>, PCollection<?>> inputs,
        @Nullable Map<TupleTag<?>, PCollection<?>> outputs) {
      this.enclosingNode = enclosingNode;
      this.transform = transform;
      this.fullName = fullName;
      this.inputs = inputs == null ? Collections.emptyMap() : inputs;
      this.outputs = outputs == null ? Collections.emptyMap() : outputs;
      this.resourceHints = transform.getResourceHints().mergeWithOuter(enclosingNode.resourceHints);
    }

    /**
     * Returns the transform associated with this transform node.
     *
     * @return {@code null} if and only if this is the root node of the graph, which has no
     *     associated transform
     */
    public @Nullable PTransform<?, ?> getTransform() {
      return transform;
    }

    /** Returns the enclosing composite transform node, or null if there is none. */
    public Node getEnclosingNode() {
      return enclosingNode;
    }

    /**
     * Adds a composite operation to the transform node.
     *
     * <p>As soon as a node is added, the transform node is considered a composite operation instead
     * of a primitive transform.
     */
    public void addComposite(Node node) {
      parts.add(node);
    }

    /**
     * Replace a sub-node of this transform node.
     *
     * <p>The existing {@link Node} must be a direct child of this {@link Node}. The replacement
     * node need not be fully specified.
     */
    public void replaceChild(Node existing, Node replacement) {
      checkNotNull(existing);
      checkNotNull(replacement);
      int existingIndex = parts.indexOf(existing);
      checkArgument(
          existingIndex >= 0,
          "Tried to replace a node %s that doesn't exist as a component of node %s",
          existing.getFullName(),
          getFullName());
      LOG.debug(
          "Replaced original node {} with replacement {} at index {}",
          existing,
          replacement,
          existingIndex);
      parts.set(existingIndex, replacement);
    }

    /**
     * Returns true if this node represents a composite transform that does not perform processing
     * of its own, but merely encapsulates a sub-pipeline (which may be empty).
     *
     * <p>Note that a node may be composite with no sub-transforms if it returns its input directly
     * extracts a component of a tuple, or other operations that occur at pipeline assembly time.
     */
    public boolean isCompositeNode() {
      return !parts.isEmpty() || isRootNode() || returnsOthersOutput();
    }

    private boolean returnsOthersOutput() {
      PTransform<?, ?> transform = getTransform();
      if (outputs != null) {
        for (PCollection<?> outputValue : outputs.values()) {
          if (!getProducer(outputValue).getTransform().equals(transform)) {
            return true;
          }
        }
      }
      return false;
    }

    public boolean isRootNode() {
      return transform == null;
    }

    public String getFullName() {
      return fullName;
    }

    /** Returns the transform input, in fully expanded form. */
    public Map<TupleTag<?>, PCollection<?>> getInputs() {
      return inputs;
    }

    /** Adds an output to the transform node. */
    private void setOutput(POutput output) {
      checkState(!finishedSpecifying);
      checkState(
          this.outputs == null, "Tried to specify more than one output for %s", getFullName());
      checkNotNull(output, "Tried to set the output of %s to null", getFullName());
      this.outputs = PValues.expandOutput(output);

      // Validate that a primitive transform produces only primitive output, and a composite
      // transform does not produce primitive output.
      Set<Node> outputProducers = new HashSet<>();
      for (PCollection<?> outputValue : PValues.expandOutput(output).values()) {
        outputProducers.add(getProducer(outputValue));
      }
      if (outputProducers.contains(this) && (!parts.isEmpty() || outputProducers.size() > 1)) {
        Set<String> otherProducerNames = new HashSet<>();
        for (Node outputProducer : outputProducers) {
          if (outputProducer != this) {
            otherProducerNames.add(outputProducer.getFullName());
          }
        }
        throw new IllegalArgumentException(
            String.format(
                "Output of composite transform [%s] contains a %s produced by it. "
                    + "Only primitive transforms are permitted to produce %ss."
                    + "%n    Outputs: %s"
                    + "%n    Other Producers: %s"
                    + "%n    Components: %s",
                getFullName(),
                PCollection.class.getSimpleName(),
                PCollection.class.getSimpleName(),
                output.expand(),
                otherProducerNames,
                parts));
      }
    }

    /**
     * Replaces each value in {@code originalToReplacement} present in this {@link Node Node's}
     * outputs with the key that maps to that value.
     *
     * <p>Outputs produced by this node that do not have an associated key in this map remain as-is.
     * Generally, these are nodes which are internal to a composite replacement; they should not
     * appear as outputs of the replacement composite.
     *
     * @param originalToReplacement A map from the outputs of the replacement {@link Node} to the
     *     original output.
     */
    void replaceOutputs(Map<PCollection<?>, ReplacementOutput> originalToReplacement) {
      checkNotNull(this.outputs, "Outputs haven't been specified for node %s yet", getFullName());
      for (Node component : this.parts) {
        // Replace the outputs of the component nodes
        component.replaceOutputs(originalToReplacement);
      }
      ImmutableMap.Builder<TupleTag<?>, PCollection<?>> newOutputsBuilder = ImmutableMap.builder();
      for (Map.Entry<TupleTag<?>, PCollection<?>> output : outputs.entrySet()) {
        ReplacementOutput mapping = originalToReplacement.get(output.getValue());
        if (mapping != null) {
          if (this.equals(producers.get(mapping.getReplacement().getValue()))) {
            // This Node produced the replacement PCollection. The structure of this if statement
            // requires the replacement transform to produce only new outputs; otherwise the
            // producers map will not be appropriately updated. TODO: investigate alternatives
            producerInput.remove(mapping.getReplacement().getValue());
            // original and replacement might be identical
            producers.remove(mapping.getReplacement().getValue());
            producers.put(mapping.getOriginal().getValue(), this);
          }
          LOG.debug(
              "Replacing output {} with original {}",
              mapping.getReplacement(),
              mapping.getOriginal());
          newOutputsBuilder.put(output.getKey(), mapping.getOriginal().getValue());
        } else {
          newOutputsBuilder.put(output);
        }
      }
      ImmutableMap<TupleTag<?>, PCollection<?>> newOutputs = newOutputsBuilder.build();
      checkState(
          outputs.size() == newOutputs.size(),
          "Number of outputs must be stable across replacement");
      this.outputs = newOutputs;
    }

    /** Returns the transform output, in expanded form. */
    public Map<TupleTag<?>, PCollection<?>> getOutputs() {
      return outputs == null ? Collections.emptyMap() : outputs;
    }

    /** Returns the {@link AppliedPTransform} representing this {@link Node}. */
    public AppliedPTransform<?, ?, ?> toAppliedPTransform(Pipeline pipeline) {
      return AppliedPTransform.of(
          getFullName(), inputs, outputs, (PTransform) getTransform(), resourceHints, pipeline);
    }

    /**
     * Visit the transform node.
     *
     * <p>The visit proceeds in the following order:
     *
     * <ul>
     *   <li>Visit all input {@link PValue PValues} returned by the flattened expansion of {@link
     *       Node#getInputs()}.
     *   <li>If the node is a composite:
     *       <ul>
     *         <li>Enter the node via {@link PipelineVisitor#enterCompositeTransform(Node)}.
     *         <li>If the result of {@link PipelineVisitor#enterCompositeTransform(Node)} was {@link
     *             CompositeBehavior#ENTER_TRANSFORM}, visit each child node of this {@link Node}.
     *         <li>Leave the node via {@link PipelineVisitor#leaveCompositeTransform(Node)}.
     *       </ul>
     *   <li>If the node is a primitive, visit it via {@link
     *       PipelineVisitor#visitPrimitiveTransform(Node)}.
     *   <li>Visit each {@link PValue} that was output by this node.
     * </ul>
     *
     * <p>Additionally, the following ordering restrictions are observed:
     *
     * <ul>
     *   <li>A {@link Node} will be visited after its enclosing node has been entered and before its
     *       enclosing node has been left
     *   <li>A {@link Node} will not be visited if any enclosing {@link Node} has returned {@link
     *       CompositeBehavior#DO_NOT_ENTER_TRANSFORM} from the call to {@link
     *       PipelineVisitor#enterCompositeTransform(Node)}.
     *   <li>A {@link PValue} will only be visited after the {@link Node} that originally produced
     *       it has been visited.
     * </ul>
     *
     * <p>Provides an ordered visit of the input values, the primitive transform (or child nodes for
     * composite transforms), then the output values.
     */
    private void visit(
        PipelineVisitor visitor,
        Set<PValue> visitedValues,
        Set<Node> visitedNodes,
        Set<Node> skippedComposites) {
      if (getEnclosingNode() != null && !visitedNodes.contains(getEnclosingNode())) {
        // Recursively enter all enclosing nodes, as appropriate.
        getEnclosingNode().visit(visitor, visitedValues, visitedNodes, skippedComposites);
      }
      // These checks occur after visiting the enclosing node to ensure that if this node has been
      // visited while visiting the enclosing node the node is not revisited, or, if an enclosing
      // Node is skipped, this node is also skipped.
      if (!visitedNodes.add(this)) {
        LOG.debug("Not revisiting previously visited node {}", this);
        return;
      } else if (childNodeOf(skippedComposites)) {
        // This node is a child of a node that has been passed over via CompositeBehavior, and
        // should also be skipped. All child nodes of a skipped composite should always be skipped.
        LOG.debug("Not revisiting Node {} which is a child of a previously passed composite", this);
        return;
      }

      if (!finishedSpecifying) {
        finishSpecifying();
      }

      if (!isRootNode()) {
        // Visit inputs.
        for (PCollection<?> inputValue : inputs.values()) {
          Node valueProducer = maybeGetProducer(inputValue);
          if (valueProducer != null) {
            if (!visitedNodes.contains(valueProducer)) {
              valueProducer.visit(visitor, visitedValues, visitedNodes, skippedComposites);
            }
            if (visitedValues.add(inputValue)) {
              LOG.debug("Visiting input value {}", inputValue);
              visitor.visitValue(inputValue, valueProducer);
            }
          }
        }
      }

      if (isCompositeNode()) {
        LOG.debug("Visiting composite node {}", this);
        PipelineVisitor.CompositeBehavior recurse = visitor.enterCompositeTransform(this);

        if (recurse.equals(CompositeBehavior.ENTER_TRANSFORM)) {
          for (Node child : parts) {
            child.visit(visitor, visitedValues, visitedNodes, skippedComposites);
          }
        } else {
          skippedComposites.add(this);
        }
        visitor.leaveCompositeTransform(this);
      } else {
        LOG.debug("Visiting primitive node {}", this);
        visitor.visitPrimitiveTransform(this);
      }

      if (!isRootNode()) {
        checkNotNull(outputs, "Outputs for non-root node %s are null", getFullName());
        // Visit outputs.
        for (PValue pValue : outputs.values()) {
          if (visitedValues.add(pValue)) {
            LOG.debug("Visiting output value {}", pValue);
            visitor.visitValue(pValue, this);
          }
        }
      }
    }

    private boolean childNodeOf(Set<Node> nodes) {
      if (isRootNode()) {
        return false;
      }
      Node parent = this.getEnclosingNode();
      while (!parent.isRootNode() && !nodes.contains(parent)) {
        parent = parent.getEnclosingNode();
      }
      return nodes.contains(parent);
    }

    /**
     * Finish specifying a transform.
     *
     * <p>All inputs are finished first, then the transform, then all outputs.
     */
    private void finishSpecifying() {
      if (finishedSpecifying) {
        return;
      }
      finishedSpecifying = true;
    }

    @Override
    public String toString() {
      if (isRootNode()) {
        return "RootNode";
      }
      return MoreObjects.toStringHelper(getClass())
          .add("fullName", fullName)
          .add("transform", transform)
          .toString();
    }
  }
}
