package org.javarosa.xpath.expr;

import org.javarosa.core.model.condition.EvaluationContext;
import org.javarosa.core.model.instance.DataInstance;
import org.javarosa.xpath.XPathArityException;
import org.javarosa.xpath.XPathTypeMismatchException;
import org.javarosa.xpath.parser.XPathSyntaxException;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import javax.xml.crypto.dsig.spec.XPathType;

/**
 * Created by amstone326 on 6/28/17.
 */

public class XPathSortFunc extends XPathFuncExpr {

    public static final String NAME = "sort";

    // since we accept 2 or 3 arguments
    private static final int EXPECTED_ARG_COUNT = -1;

    public XPathSortFunc() {
        name = NAME;
        expectedArgCount = EXPECTED_ARG_COUNT;
    }

    public XPathSortFunc(XPathExpression[] args) throws XPathSyntaxException {
        super(NAME, args, EXPECTED_ARG_COUNT, true);
    }

    @Override
    protected void validateArgCount() throws XPathSyntaxException {
        if (args.length != 2 && args.length != 3) {
            throw new XPathArityException(name, "2 or 3 arguments", args.length);
        }
    }

    @Override
    protected Object evalBody(DataInstance model, EvaluationContext evalContext, Object[] evaluatedArgs) {
        List<String> sortedList;
        if (evaluatedArgs.length == 2) {
            sortedList =
                    sortSingleList(FunctionUtils.toString(evaluatedArgs[0]),
                            FunctionUtils.toBoolean(evaluatedArgs[1]));
        } else {
            sortedList =
                    sortListByOtherList(FunctionUtils.toString(evaluatedArgs[0]),
                            FunctionUtils.toString(evaluatedArgs[1]),
                            FunctionUtils.toBoolean(evaluatedArgs[2]));
        }
        return listToString(sortedList);
    }

    private List<String> sortSingleList(String spaceSeparatedString, final boolean ascending) {
        List<String> items = stringToList(spaceSeparatedString);
        Collections.sort(items, new Comparator<String>() {
            @Override
            public int compare(String s1, String s2) {
                return (ascending ? 1 : -1) * s1.compareTo(s2);
            }
        });
        return items;
    }

    private List<String> sortListByOtherList(String s1, String s2, boolean ascending) {
        List<String> targetListItems = stringToList(s1);
        List<String> comparisonListItems = stringToList(s2);

        if (targetListItems.size() != comparisonListItems.size()) {
            throw new XPathTypeMismatchException("Length of lists passed to sort() must match");
        }

        Map<String, Object> stringMapping =
                createMappingFromComparisonToTarget(comparisonListItems, targetListItems);

        List<String> sortedComparisonList = sortSingleList(s2, ascending);
        List<String> sortedTargetList = new ArrayList<>();
        String previousComparisonString = "";
        for (int i = 0; i < sortedComparisonList.size(); i++) {
            String stringInSortedList = sortedComparisonList.get(i);
            if (stringInSortedList.equals(previousComparisonString)) {
                // Means we already grabbed all the target strings corresponding to this reference string
                continue;
            }
            Object mappedStrings = stringMapping.get(stringInSortedList);
            if (mappedStrings instanceof String) {
                sortedTargetList.add((String)mappedStrings);
            } else {
                sortedTargetList.addAll((List<String>)mappedStrings);
            }
            previousComparisonString = stringInSortedList;
        }

        return sortedTargetList;
    }

    private static Map<String, Object> createMappingFromComparisonToTarget(List<String> comparisonListItems,
                                                                               List<String> targetListItems) {
        Map<String, Object> stringMapping = new HashMap<>();
        for (int i = 0; i < comparisonListItems.size(); i++) {
            String comparisonString = comparisonListItems.get(i);
            String targetString = targetListItems.get(i);
            if (stringMapping.containsKey(comparisonString)) {
                Object alreadyMapped = stringMapping.get(comparisonString);
                if (alreadyMapped instanceof String) {
                    List<String> multipleValues = new ArrayList<>();
                    multipleValues.add((String)alreadyMapped);
                    multipleValues.add(targetString);
                    stringMapping.put(comparisonString, multipleValues);
                } else {
                    ((List<String>)alreadyMapped).add(targetString);
                }
            } else {
                stringMapping.put(comparisonString, targetString);
            }
        }
        return stringMapping;
    }

    private static List<String> stringToList(String s) {
        return Arrays.asList(s.split(" "));
    }

    private static String listToString(List<String> list) {
        StringBuilder sb = new StringBuilder();
        for (String s : list) {
            sb.append(s + " ");
        }
        return sb.toString().substring(0, sb.length()-1);
    }

}

@@ -2748,38 +2748,44 @@
     Map%3CString, 
-Object
+List%3CString%3E
 %3E stringMapping 
@@ -3412,137 +3412,8 @@
   %7D%0A
-            Object mappedStrings = stringMapping.get(stringInSortedList);%0A            if (mappedStrings instanceof String) %7B%0A    
     
@@ -3444,137 +3444,51 @@
 .add
-((S
+All(s
 tring
-)mappedStrings);%0A            %7D else %7B%0A                sortedTargetList.addAll((List%3CString%3E)mappedStrings);%0A            %7D
+Mapping.get(stringInSortedList));
 %0A   
@@ -3625,22 +3625,28 @@
 String, 
-Object
+List%3CString%3E
 %3E create
@@ -3776,36 +3776,32 @@
                 
-    
 List%3CString%3E tar
@@ -3836,22 +3836,28 @@
 String, 
-Object
+List%3CString%3E
 %3E string
@@ -4062,32 +4062,79 @@
 stItems.get(i);%0A
+            List%3CString%3E correspondingStrings;%0A
             if (
@@ -4200,28 +4200,28 @@
     
-Object alreadyMapped
+correspondingStrings
  = s
@@ -4272,408 +4272,150 @@
     
-    if (alreadyMapped instanceof String) %7B%0A                    List%3CString%3E multipleValues = new ArrayList%3C%3E();%0A                    multipleValues.add((String)alreadyMapped);%0A                    multipleValues.add(targetString);%0A                    stringMapping.put(comparisonString, multipleValues);%0A                %7D else %7B%0A                    ((List%3CString%3E)alreadyMapped).add(target
+%7D else %7B%0A                correspondingStrings = new ArrayList%3C%3E();%0A                stringMapping.put(comparisonString, corresponding
 String
+s
 );%0A
-    
     
@@ -4440,97 +4440,47 @@
     
-%7D else %7B%0A                stringMapping.put(comparisonString, targetString);%0A            %7D
+correspondingStrings.add(targetString);
 %0A   

simplify to always using a list