package io.burt.jmespath.function;

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

import io.burt.jmespath.Adapter;
import io.burt.jmespath.Expression;
import io.burt.jmespath.JmesPathType;

public class SortByFunction extends BaseFunction {
  public SortByFunction() {
    super(
      ArgumentConstraints.arrayOf(ArgumentConstraints.typeOf(JmesPathType.OBJECT)),
      ArgumentConstraints.expression()
    );
  }

  @Override
  protected <T> T callFunction(final Adapter<T> runtime, List<FunctionArgument<T>> arguments) {
    List<T> elementsList = runtime.toList(arguments.get(0).value());
    Expression<T> expression = arguments.get(1).expression();
    Iterator<T> elements = elementsList.iterator();
    if (elements.hasNext()) {
      List<Pair<T>> pairs = new ArrayList<>(elementsList.size());
      T element = elements.next();
      T transformedElement = expression.search(element);
      boolean expectNumbers = true;
      JmesPathType elementType = runtime.typeOf(transformedElement);
      if (elementType == JmesPathType.STRING) {
        expectNumbers = false;
      } else if (elementType != JmesPathType.NUMBER) {
        throw new ArgumentTypeException(name(), "number or string", elementType.toString());
      }
      pairs.add(new Pair<T>(transformedElement, element));
      while (elements.hasNext()) {
        element = elements.next();
        transformedElement = expression.search(element);
        elementType = runtime.typeOf(transformedElement);
        if (expectNumbers && elementType != JmesPathType.NUMBER) {
          throw new ArgumentTypeException(name(), "number", elementType.toString());
        } else if (!expectNumbers && elementType != JmesPathType.STRING) {
          throw new ArgumentTypeException(name(), "string", elementType.toString());
        }
        pairs.add(new Pair<T>(transformedElement, element));
      }
      Collections.sort(pairs, new Comparator<Pair<T>>() {
        @Override
        public int compare(Pair<T> a, Pair<T> b) {
          return runtime.compare(a.transformedElement, b.transformedElement);
        }
      });
      List<T> sorted = new ArrayList<>(pairs.size());
      for (Pair<T> pair : pairs) {
        sorted.add(pair.element);
      }
      return runtime.createArray(sorted);
    } else {
      return runtime.createArray(new ArrayList<T>());
    }
  }

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

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

@@ -1013,16 +1013,595 @@
 s = 
-true;%0A  
+expectNumbers(runtime, transformedElement );%0A      pairs.add(new Pair%3CT%3E(transformedElement, element));%0A      while (elements.hasNext()) %7B%0A        element = elements.next();%0A        transformedElement = expression.search(element);%0A        checkType(runtime, transformedElement, expectNumbers);%0A        pairs.add(new Pair%3CT%3E(transformedElement, element));%0A      %7D%0A      return runtime.createArray(sortAndFlatten(runtime, pairs));%0A    %7D else %7B%0A      return runtime.createArray(new ArrayList%3CT%3E());%0A    %7D%0A  %7D%0A%0A  private %3CT%3E boolean expectNumbers(Adapter%3CT%3E runtime, T transformedElement) %7B%0A
     
@@ -1663,26 +1663,24 @@
 ement);%0A    
-  
 if (elementT
@@ -1719,33 +1719,22 @@
     
-  expectNumbers =
+return
  false;%0A
     
@@ -1729,18 +1729,16 @@
  false;%0A
-  
     %7D el
@@ -1778,26 +1778,24 @@
 e.NUMBER) %7B%0A
-  
       throw 
@@ -1869,26 +1869,24 @@
 oString());%0A
-  
     %7D%0A      
@@ -1887,197 +1887,138 @@
     
-  pairs.add(new Pair%3CT%3E(transformedElement, element));%0A      while (elements.hasNext()) %7B%0A        element = elements.next();%0A        transformedElement = expression.search(element);%0A       
+return true;%0A  %7D%0A%0A  private %3CT%3E void checkType(Adapter%3CT%3E runtime, T transformedElement, boolean expectNumbers) %7B%0A    JmesPathType
  ele
@@ -2060,28 +2060,24 @@
 edElement);%0A
-    
     if (expe
@@ -2125,36 +2125,32 @@
 NUMBER) %7B%0A      
-  
-  
 throw new Argume
@@ -2204,28 +2204,24 @@
 oString());%0A
-    
     %7D else i
@@ -2285,20 +2285,16 @@
 %7B%0A      
-  
-  
 throw ne
@@ -2368,85 +2368,101 @@
     
-    %7D%0A        pairs.add(new Pair%3CT%3E(transformedElement, element));%0A      %7D%0A  
+%7D%0A  %7D%0A%0A  private %3CT%3E List%3CT%3E sortAndFlatten(final Adapter%3CT%3E runtime, List%3CPair%3CT%3E%3E pairs) %7B%0A
     
@@ -2509,26 +2509,24 @@
 air%3CT%3E%3E() %7B%0A
-  
       @Overr
@@ -2535,18 +2535,16 @@
 e%0A      
-  
 public i
@@ -2574,26 +2574,24 @@
 air%3CT%3E b) %7B%0A
-  
         retu
@@ -2656,30 +2656,26 @@
 ent);%0A      
-  %7D%0A  
+%7D%0A
     %7D);%0A    
@@ -2670,26 +2670,24 @@
     %7D);%0A    
-  
 List%3CT%3E sort
@@ -2718,26 +2718,24 @@
 rs.size());%0A
-  
     for (Pai
@@ -2761,18 +2761,16 @@
 %7B%0A      
-  
 sorted.a
@@ -2779,34 +2779,32 @@
 (pair.element);%0A
-  
     %7D%0A      retu
@@ -2801,118 +2801,22 @@
     
-  
 return 
-runtime.createArray(sorted);%0A    %7D else %7B%0A      return runtime.createArray(new ArrayList%3CT%3E());%0A    %7D
+sorted;
 %0A  %7D

Decrease complexity in SortByFunction