package io.burt.jmespath.function;

import java.util.Collections;
import java.util.List;
import java.util.ArrayList;
import java.util.Comparator;

import io.burt.jmespath.Adapter;

public class SortByFunction extends ComparingFunction {
  @Override
  protected <T> ComparingFunction.Aggregator<T> createAggregator(Adapter<T> runtime, T element, T elementValue) {
    return new SortingAggregator<T>(runtime, element, elementValue);
  }

  @Override
  protected <T> T createNullValue(Adapter<T> runtime) {
    return runtime.createArray(new ArrayList<T>());
  }

  private class SortingAggregator<V> extends ComparingFunction.Aggregator<V> {
    private List<Pair<V>> pairs;

    public SortingAggregator(Adapter<V> runtime, V initialElement, V initialValue) {
      super(runtime);
      this.pairs = new ArrayList<>();
      this.pairs.add(new Pair<V>(initialElement, initialValue));
    }

    protected void aggregate(V candidate, V candidateValue) {
      pairs.add(new Pair<V>(candidate, candidateValue));
    }
  
    protected V result() {
      return runtime.createArray(sortAndFlatten(runtime, pairs));
    }
  }

  private <T> List<T> sortAndFlatten(final Adapter<T> runtime, List<Pair<T>> pairs) {
    Collections.sort(pairs, new Comparator<Pair<T>>() {
      @Override
      public int compare(Pair<T> a, Pair<T> b) {
        return runtime.compare(a.elementValue, b.elementValue);
      }
    });
    List<T> sorted = new ArrayList<>(pairs.size());
    for (Pair<T> pair : pairs) {
      sorted.add(pair.element);
    }
    return sorted;
  }

  private static class Pair<U> {
    public final U element;
    public final U elementValue;

    public Pair(U element, U elementValue) {
      this.element = element;
      this.elementValue = elementValue;
    }
  }
}

@@ -1112,21 +1112,19 @@
 ;%0A    %7D%0A
+%0A
   
-%7D%0A%0A
   privat
@@ -1205,16 +1205,18 @@
 ) %7B%0A    
+  
 Collecti
@@ -1255,24 +1255,26 @@
 air%3CT%3E%3E() %7B%0A
+  
       @Overr
@@ -1279,24 +1279,26 @@
 rride%0A      
+  
 public int c
@@ -1324,24 +1324,26 @@
 air%3CT%3E b) %7B%0A
+  
         retu
@@ -1396,26 +1396,30 @@
 lue);%0A      
-%7D%0A
+  %7D%0A  
     %7D);%0A    
@@ -1418,16 +1418,18 @@
 %7D);%0A    
+  
 List%3CT%3E 
@@ -1464,24 +1464,26 @@
 rs.size());%0A
+  
     for (Pai
@@ -1509,16 +1509,18 @@
 %7B%0A      
+  
 sorted.a
@@ -1537,16 +1537,18 @@
 ement);%0A
+  
     %7D%0A  
@@ -1545,24 +1545,26 @@
       %7D%0A    
+  
 return sorte
@@ -1566,16 +1566,22 @@
 sorted;%0A
+    %7D%0A
   %7D%0A%0A  p

Move SortByFunction#sortAndFlatten into the aggregator class