package org.batfish.datamodel.routing_policy.statement;

import com.fasterxml.jackson.annotation.JsonCreator;
import com.google.common.collect.ImmutableSortedSet;
import java.util.List;
import java.util.SortedSet;
import org.batfish.datamodel.BgpRoute;
import org.batfish.datamodel.routing_policy.Environment;
import org.batfish.datamodel.routing_policy.Result;
import org.batfish.datamodel.routing_policy.expr.AsPathListExpr;

public class PrependAsPath extends Statement {

  /** */
  private static final long serialVersionUID = 1L;

  private AsPathListExpr _expr;

  @JsonCreator
  private PrependAsPath() {}

  public PrependAsPath(AsPathListExpr expr) {
    _expr = expr;
  }

  @Override
  public boolean equals(Object obj) {
    if (this == obj) {
      return true;
    }
    if (obj == null) {
      return false;
    }
    if (getClass() != obj.getClass()) {
      return false;
    }
    PrependAsPath other = (PrependAsPath) obj;
    if (_expr == null) {
      if (other._expr != null) {
        return false;
      }
    } else if (!_expr.equals(other._expr)) {
      return false;
    }
    return true;
  }

  @Override
  public Result execute(Environment environment) {
    List<Integer> toPrepend = _expr.evaluate(environment);
    BgpRoute.Builder bgpRouteBuilder = (BgpRoute.Builder) environment.getOutputRoute();
    List<SortedSet<Integer>> asPath = bgpRouteBuilder.getAsPath();
    List<SortedSet<Integer>> iAsPath = null;
    if (environment.getWriteToIntermediateBgpAttributes()) {
      BgpRoute.Builder ir = environment.getIntermediateBgpAttributes();
      iAsPath = ir.getAsPath();
    }
    for (Integer as : toPrepend) {
      SortedSet<Integer> asSet = ImmutableSortedSet.of(as);
      asPath.add(0, asSet);
      if (iAsPath != null) {
        iAsPath.add(0, asSet);
      }
    }
    Result result = new Result();
    return result;
  }

  public AsPathListExpr getExpr() {
    return _expr;
  }

  @Override
  public int hashCode() {
    final int prime = 31;
    int result = 1;
    result = prime * result + ((_expr == null) ? 0 : _expr.hashCode());
    return result;
  }

  public void setExpr(AsPathListExpr expr) {
    _expr = expr;
  }
}

@@ -203,24 +203,60 @@
 .SortedSet;%0A
+import java.util.stream.Collectors;%0A
 import org.b
@@ -1272,24 +1272,152 @@
 vironment);%0A
+    List%3CSortedSet%3CInteger%3E%3E newAsPaths =%0A        toPrepend.stream().map(ImmutableSortedSet::of).collect(Collectors.toList());%0A%0A
     BgpRoute
@@ -1567,48 +1567,38 @@
     
-List%3CSortedSet%3CInteger%3E%3E iAsPath = null;
+asPath.addAll(0, newAsPaths);%0A
 %0A   
@@ -1729,24 +1729,49 @@
 es();%0A      
+List%3CSortedSet%3CInteger%3E%3E 
 iAsPath = ir
@@ -1788,174 +1788,14 @@
 ();%0A
-    %7D%0A    for (Integer as : toPrepend) %7B%0A      SortedSet%3CInteger%3E asSet = ImmutableSortedSet.of(as);%0A      asPath.add(0, asSet);%0A      if (iAsPath != null) %7B%0A 
       
- 
 iAsP
@@ -1805,33 +1805,34 @@
 .add
+All
 (0, 
-asSet
+newAsPaths
 );%0A
-  
     %7D%0A
-    %7D
 %0A   

optimize algorithm