package org.batfish.common.traceroute;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;

import com.google.common.collect.ImmutableList;
import java.util.List;
import java.util.Stack;
import java.util.stream.Stream;
import javax.annotation.Nullable;
import org.batfish.datamodel.Flow;
import org.batfish.datamodel.FlowDisposition;
import org.batfish.datamodel.flow.FirewallSessionTraceInfo;
import org.batfish.datamodel.flow.Hop;
import org.batfish.datamodel.flow.Trace;
import org.batfish.datamodel.flow.TraceAndReverseFlow;

/**
 * A DAG representation of a collection of {@link TraceAndReverseFlow traces}. Every path through
 * the DAG from a root to a leaf is a valid {@link TraceAndReverseFlow trace}, i.e. one that {@link
 * org.batfish.common.plugin.TracerouteEngine} would create.
 */
public final class TraceDagImpl implements TraceDag {
  /** A node in the DAG contains everything needed to reconstruct traces through the node. */
  public static final class Node {
    private final Hop _hop;
    private final @Nullable FirewallSessionTraceInfo _firewallSessionTraceInfo;
    private final @Nullable FlowDisposition _flowDisposition;
    private final @Nullable Flow _returnFlow;
    private final List<Integer> _successors;

    public Node(
        Hop hop,
        @Nullable FirewallSessionTraceInfo firewallSessionTraceInfo,
        @Nullable FlowDisposition flowDisposition,
        @Nullable Flow returnFlow,
        List<Integer> successors) {
      checkArgument(
          (flowDisposition != null && flowDisposition.isSuccessful()) == (returnFlow != null),
          "returnFlow is present iff disposition is successful");
      _hop = hop;
      _firewallSessionTraceInfo = firewallSessionTraceInfo;
      _flowDisposition = flowDisposition;
      _returnFlow = returnFlow;
      _successors = ImmutableList.copyOf(successors);
    }
  }

  private final List<Node> _nodes;
  private final List<Integer> _rootIds;

  public TraceDagImpl(List<Node> nodes, List<Integer> rootIds) {
    _nodes = nodes;
    _rootIds = rootIds;
  }

  @Override
  public int size() {
    return new SizeComputer().size();
  }

  /** Efficiently compute the number of traces in the DAG. */
  private class SizeComputer {
    private final Integer[] _cache;

    SizeComputer() {
      _cache = new Integer[_nodes.size()];
    }

    int size(int nodeId) {
      @Nullable Integer cachedSize = _cache[nodeId];
      if (cachedSize != null) {
        return cachedSize;
      }
      Node node = _nodes.get(nodeId);
      if (node._successors.isEmpty()) {
        // leaf
        _cache[nodeId] = 1;
      } else {
        _cache[nodeId] = node._successors.stream().mapToInt(this::size).sum();
      }
      return _cache[nodeId];
    }

    int size() {
      return _rootIds.stream().mapToInt(this::size).sum();
    }
  }

  private Stream<TraceAndReverseFlow> getTraces(
      List<Hop> hopsInput, List<FirewallSessionTraceInfo> sessionsInput, int rootId) {
    Node node = _nodes.get(rootId);

    List<Hop> hops =
        ImmutableList.<Hop>builderWithExpectedSize(hopsInput.size() + 1)
            .addAll(hopsInput)
            .add(node._hop)
            .build();

    List<FirewallSessionTraceInfo> sessions;
    FirewallSessionTraceInfo firewallSessionTraceInfo = node._firewallSessionTraceInfo;
    if (firewallSessionTraceInfo != null) {
      sessions =
          ImmutableList.<FirewallSessionTraceInfo>builderWithExpectedSize(sessionsInput.size() + 1)
              .addAll(sessionsInput)
              .add(firewallSessionTraceInfo)
              .build();
    } else {
      sessions = sessionsInput;
    }

    List<Integer> successors = node._successors;
    if (successors.isEmpty()) {
      FlowDisposition disposition =
          checkNotNull(node._flowDisposition, "failed to determine disposition from hop");
      Trace trace = new Trace(disposition, hops);
      return Stream.of(new TraceAndReverseFlow(trace, node._returnFlow, sessions));
    } else {
      return successors.stream().flatMap(successorId -> getTraces(hops, sessions, successorId));
    }
  }

  @Override
  public int countEdges() {
    return _nodes.stream().map(node -> node._successors).mapToInt(List::size).sum();
  }

  @Override
  public int countNodes() {
    return _nodes.size();
  }

  private Stream<TraceAndReverseFlow> getTraces(int rootId) {
    return getTraces(ImmutableList.of(), new Stack<>(), rootId);
  }

  @Override
  public Stream<TraceAndReverseFlow> getTraces() {
    return _rootIds.stream().flatMap(this::getTraces);
  }
}

@@ -292,16 +292,49 @@
 Stream;%0A
+import javax.annotation.Nonnull;%0A
 import j
@@ -359,16 +359,16 @@
 llable;%0A
-
 import o
@@ -1104,16 +1104,25 @@
 e final 
+@Nonnull 
 Hop _hop
@@ -1760,24 +1760,171 @@
 ccessful%22);%0A
+      checkArgument(%0A          (flowDisposition != null) == successors.isEmpty(),%0A          %22flowDisposition is present iff successors is empty%22);%0A
       _hop =
@@ -2114,32 +2114,478 @@
 ccessors);%0A    %7D
+%0A%0A    public @Nonnull Hop getHop() %7B%0A      return _hop;%0A    %7D%0A%0A    public @Nullable FirewallSessionTraceInfo getFirewallSessionTraceInfo() %7B%0A      return _firewallSessionTraceInfo;%0A    %7D%0A%0A    public @Nullable FlowDisposition getFlowDisposition() %7B%0A      return _flowDisposition;%0A    %7D%0A%0A    public @Nullable Flow getReturnFlow() %7B%0A      return _returnFlow;%0A    %7D%0A%0A    public @Nonnull List%3CInteger%3E getSuccessors() %7B%0A      return _successors;%0A    %7D
 %0A  %7D%0A%0A  private 
@@ -2735,16 +2735,16 @@
  nodes;%0A
-
     _roo
@@ -2760,24 +2760,90 @@
 otIds;%0A  %7D%0A%0A
+  public @Nonnull List%3CNode%3E getNodes() %7B%0A    return _nodes;%0A  %7D%0A%0A
   @Override%0A

add invariant, expose node fields (#6681)