/*
 * Copyright 2018 Scytl Secure Electronic Voting SA
 *
 * All rights reserved
 *
 * See our extended copyright notice in *file 'Copyright.txt' which is part of this source code package
 */
package com.scytl.products.ov.validation.commons.writeins;

import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Base64;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

import com.scytl.products.ov.validation.commons.ValidationCommonsException;

import org.apache.commons.lang.StringUtils;

import static com.scytl.products.ov.commons.util.MathUtils.isqrt;

/**
 * Permits the encoding and validation of a list of strings as a list of integer values, and also permits the decoding
 * of such lists of integers back to the original list of strings. Validations are again performed at the moment of
 * decoding (to detect if the encoding is valid). The validations that are performed ensure that the input strings match
 * an expected pattern. The pattern that the input strings must match is the following:
 * <P>
 * [Representation as a string][Separator character][Actual data]
 * <P>
 * Additionally, the validations ensure that the representations that appear in each string also appear in a list of
 * valid representations that are provided as arguments at the moment of encoding and decoding. Each of the
 * representations in the provided list of valid representations must appear exactly once in the list of strings to be
 * encoded.
 * <P>
 * The number of elements in the generated encoding (in other words, the size of the encoding) will match exactly the
 * size parameter that is specified at the moment of encoding.
 * <P>
 * The integers that compose the encoding will be members of the group defined by the p parameter, specified at the
 * moment of encoding.
 * <P>
 * This encoder can only encode strings that are only composed of characters from a particular alphabet. The first
 * character in the alphabet is used as a special separator character. Each of the strings to be encoded should contain
 * within it exactly one of the separator characters. The format of each of the strings must match the following format:
 * The approach that this encoder follows is to initially encode the input data using the minimum of integer values, and
 * then later add padding to the encoding so that is matches the desired size.
 * <P>
 * To perform the encoding, each character in the input data is mapped to an integer value (the index of that character
 * in the supported alphabet). A certain number of those integers are then concatenated together (a separator character
 * is inserted between each separate string), finally that value is converted to a group element (by raising it to the
 * power of two).
 * <P>
 * The number of characters that can be included in a single group element depends on the size (in bits) of the 'p'
 * parameter of the group. Given a 'p' parameter of a particular size, the following shows how many characters can be
 * encoded in a single group element:
 * <ul>
 * <li>P size = 1024, number of characters that can be encoded=51</li>
 * <li>P size = 2048, number of characters that can be encoded=102</li>
 * <li>P size = 3072, number of characters that can be encoded=153</li>
 * </ul>
 * A factor that should be taken into account is the number of separate strings that need to be encoded, because one
 * character of the available characters is consumed by the separator characters between each string. For example, 2
 * strings of length 30 characters equate to 61 characters. 3 strings of length 30 characters equate to 92 characters,
 * etc.
 */
public class WriteinEncoder {

    private static final BigInteger PADDING_INTEGER_VALUE = BigInteger.valueOf(2);

    private static final int NUM_CHARS_PER_ENCODED_CHARACTER = 3;

    private static final String PADDING_CHARACTER = "0";

    private static final String TWO_PADDING_CHARACTERS = PADDING_CHARACTER + PADDING_CHARACTER;

    // this value was determined empirically
    private static final int FACTOR_TO_DETERMINE_NUM_CHARS_TO_ENCODE_IN_EACH_GROUP_ELEMENT = 20;

    private static final int OFFSET = 2;

    private final String separator;

    private final String quotedSeparator;

    private final String alphabet;

    private final Pattern pattern;
    
       
    public WriteinEncoder(String alphabetBase64) {

        if ((alphabetBase64 == null) || (alphabetBase64.isEmpty())) {
            throw new IllegalArgumentException("The received alphabet was not an initialized non-empty string");
        }

        this.alphabet = new String(Base64.getDecoder().decode(alphabetBase64), StandardCharsets.UTF_8);
        this.separator = alphabet.substring(0, 1);
        this.quotedSeparator = Pattern.quote(this.separator);
		this.pattern = Pattern.compile("(\\d)+" + this.quotedSeparator + "[^" + this.quotedSeparator + "]*+("
				+ this.quotedSeparator + "(\\d)+" + this.quotedSeparator + "[^" + this.quotedSeparator + "]*)*");
	}

    /**
     * Encodes a list of strings as a list of integer values that are elements of the group defined by the p parameter.
     * 
     * @param strings
     *            the list of strings to be encoded.
     * @param p
     *            the p parameter of the group.
     * @param numDesiredElementsInEncoding
     *            the desired size of the encoding
     * @param validWriteInRepresentations
     *            the list of valid integer values that can be associated with a writein string.
     * @return the encoding, represented as a list of integer values that are members of the group defined by the p
     *         parameter.
     */
    public List<BigInteger> encode(List<String> strings, BigInteger p, BigInteger q, int numDesiredElementsInEncoding,
            List<String> validWriteInRepresentations) {

        if (numDesiredElementsInEncoding < 1) {
            throw new IllegalArgumentException("The number of desired elements was " + numDesiredElementsInEncoding);
        }

        if ((strings == null) || (strings.isEmpty())) {

            if (validWriteInRepresentations == null || validWriteInRepresentations.isEmpty()) {

                List<BigInteger> encoding = new ArrayList<BigInteger>();
                for (int i = 0; i < numDesiredElementsInEncoding; i++) {
                    encoding.add(PADDING_INTEGER_VALUE);
                }
                return encoding;

            } else {
                throw new IllegalArgumentException(
                    "Exception during encoding - there are valid representations but no write-ins");
            }
        }

        performBasicValidationsOnEncodeInputs(strings, p, q);

        confirmThatRepsAreUniqueAndAppearExacltOnceInListOfValidReps(strings, validWriteInRepresentations);

        List<BigInteger> encoding = performEncoding(strings, p, q);

        int numElementsInRealEncoding = encoding.size();
        if (numElementsInRealEncoding > numDesiredElementsInEncoding) {
            throw new ValidationCommonsException("The encoding has more elements than the desired number. Desired number: "
                + numDesiredElementsInEncoding + ", actual number: " + numElementsInRealEncoding);
        }

        int numRequiredPaddingElements = numDesiredElementsInEncoding - numElementsInRealEncoding;
        for (int i = 0; i < numRequiredPaddingElements; i++) {
            encoding.add(PADDING_INTEGER_VALUE);
        }

        return encoding;
    }

    private void confirmThatRepsAreUniqueAndAppearExacltOnceInListOfValidReps(List<String> strings,
            List<String> validWriteInRepresentations) {

        Set<String> validWriteInRepresentationsAsSet = new HashSet<>(validWriteInRepresentations);
        Set<String> setToConfirmRepsUnique = new HashSet<>();

        for (int i = 0; i < strings.size(); i++) {

            String[] parts = strings.get(i).split(separator);
            String rep = parts[0];

            if (!setToConfirmRepsUnique.add(rep)) {
                throw new IllegalArgumentException("Exception during encoding - repeated representation: " + rep);
            }

            if (!validWriteInRepresentationsAsSet.remove(rep)) {
                throw new IllegalArgumentException(
                    "Exception during encoding - invalid representation found in encoded write-in: " + rep);
            }
        }
        if (!validWriteInRepresentationsAsSet.isEmpty()) {
            throw new IllegalArgumentException(
                "Exception during encoding - representations not used: " + validWriteInRepresentationsAsSet);
        }
    }

    /**
     * Decode the received encoding to a list of strings.
     * 
     * @param encoding
     *            the encoding to be decoded.
     * @param p
     *            the p parameter of the group. This must have the same value as was used in the encode operation.
     * @param validWriteInRepresentations
     *            the list of valid integer values that can be associated with a writein string.
     * @return the decoded data, as a list of strings.
     */
    public List<String> decode(List<BigInteger> encoding, BigInteger p, List<String> validWriteInRepresentations) {

        performBasicValidationsOnDecodeInputs(p);

        List<BigInteger> listWithoutPadding = encoding.subList(0, findNumberOfNonPaddingElements(encoding));

        return performDecoding(listWithoutPadding, p, validWriteInRepresentations);
    }

    /**
     * @return the alphabet supported by this encoder.
     */
    public String getAlphabet() {
        return alphabet;
    }

    private int findNumberOfNonPaddingElements(List<BigInteger> encoding) {

        int numNonPaddingElements = 0;

        for (BigInteger e : encoding) {

            if (!(PADDING_INTEGER_VALUE.equals(e))) {
                numNonPaddingElements++;
            }
        }

        return numNonPaddingElements;
    }

    private List<BigInteger> performEncoding(List<String> strings, BigInteger p, BigInteger q) {

        int numCharsThatCanBeEncodedInEachGroupElement =
            p.bitLength() / FACTOR_TO_DETERMINE_NUM_CHARS_TO_ENCODE_IN_EACH_GROUP_ELEMENT;

        String concatenatedWithSeparators = String.join(separator, strings);

        int concatenatedWithSeparatorsLength = concatenatedWithSeparators.length();
        List<BigInteger> returnList = new ArrayList<BigInteger>();
        int numCharsAddedToThisGroupElement = 0;
        int totalNumCharsProcessed = 0;

        int lengthOfStringThatEncodesMaximumNumberOfCharactersThatGroupElementCanHold =
            numCharsThatCanBeEncodedInEachGroupElement * NUM_CHARS_PER_ENCODED_CHARACTER;

        int optimalInitialStringBuilderLength =
            Math.min(lengthOfStringThatEncodesMaximumNumberOfCharactersThatGroupElementCanHold,
                concatenatedWithSeparatorsLength * NUM_CHARS_PER_ENCODED_CHARACTER);
        StringBuilder sb = new StringBuilder(optimalInitialStringBuilderLength);

        for (char character : concatenatedWithSeparators.toCharArray()) {

            int indexOfCharacterInAlphabet = alphabet.indexOf(character) + 1;
            if (indexOfCharacterInAlphabet < 1) {
                throw new IllegalArgumentException("Invalid character for writein encoding: " + character);
            }

            indexOfCharacterInAlphabet += OFFSET;
            sb.append(getStringRepresentationAndAddPaddingCharsIfNecessary(indexOfCharacterInAlphabet));
            numCharsAddedToThisGroupElement++;
            totalNumCharsProcessed++;

            if (numCharsAddedToThisGroupElement == numCharsThatCanBeEncodedInEachGroupElement) {

                returnList.add(squareToEnsureIsGroupElement(new BigInteger(sb.toString()), p, q));
                optimalInitialStringBuilderLength =
                    Math.min(lengthOfStringThatEncodesMaximumNumberOfCharactersThatGroupElementCanHold,
                        (concatenatedWithSeparatorsLength - totalNumCharsProcessed) * NUM_CHARS_PER_ENCODED_CHARACTER);
                sb = new StringBuilder(optimalInitialStringBuilderLength);
                numCharsAddedToThisGroupElement = 0;
            }
        }
        if (numCharsAddedToThisGroupElement != 0) {
            returnList.add(squareToEnsureIsGroupElement(new BigInteger(sb.toString()), p, q));
        }

        return returnList;
    }

    private List<String> performDecoding(List<BigInteger> encoding, BigInteger p,
            List<String> usableWriteInRepresentations) {

        if ((encoding == null) || (encoding.isEmpty())) {

            if (usableWriteInRepresentations.isEmpty()) {
                return new ArrayList<>();
            } else {
                throw new IllegalArgumentException(
                    "Exception during decoding - there are valid representations but no write-ins");
            }
        }

        StringBuilder sb = new StringBuilder(calculateOptimalInitialStringBuilderLength(encoding));

        for (BigInteger groupElement : encoding) {

            String originalIntegerValueAsString = isqrt(groupElement.mod(p)).toString();
            int numPaddingCharsNeeded = determineNumPaddingCharsNeeded(originalIntegerValueAsString);

            for (int i = 0; i < numPaddingCharsNeeded; i++) {
                sb.append(PADDING_CHARACTER);
            }
            sb.append(originalIntegerValueAsString);
        }

        String decode = null;
        try{
            decode = decode(sb.toString());
        } catch (RuntimeException e){
            throw new IllegalArgumentException("The write-ins cannot be decoded to a valid string", e);
        }
        if (!pattern.matcher(decode).matches()) {
            throw new IllegalArgumentException("The decoded write-ins contain invalid character sequence");
        }
        String[] bothPartsOfAllStrings = decode.split(quotedSeparator, -1);
        return handleSeparatorsWithinEachString(bothPartsOfAllStrings, usableWriteInRepresentations);
    }

    private List<String> handleSeparatorsWithinEachString(String[] originalIndividualStrings,
            List<String> usableWriteInRepresentations) {

        confirmAllElementsUnique(originalIndividualStrings);

        Set<String> usableWriteInRepresentationsAsSet = new HashSet<>(usableWriteInRepresentations);
        ArrayList<String> listItems = new ArrayList<String>();
        for (int i = 0; i < originalIndividualStrings.length; i = i + 2) {
            
            if (StringUtils.isEmpty(originalIndividualStrings[i + 1])) {
                throw new IllegalArgumentException("Exception during decoding - empty write-in is not valid");
            }
            
            if (usableWriteInRepresentationsAsSet.remove(originalIndividualStrings[i])) {
                listItems.add(originalIndividualStrings[i] + separator + originalIndividualStrings[i + 1]);
            } else {
                throw new IllegalArgumentException(
                    "Exception during decoding - invalid representation found in encoded write-in: "
                        + originalIndividualStrings[i]);
            }
        }
        if (!usableWriteInRepresentationsAsSet.isEmpty()) {
            throw new IllegalArgumentException(
                "Exception during decoding - representations not used: " + usableWriteInRepresentationsAsSet);
        }
        return listItems;
    }

    private void confirmAllElementsUnique(String[] originalIndividualStrings) {

        Set<String> set = new HashSet<>();
        for (int i = 0; i < originalIndividualStrings.length; i = i + 2) {
            if (!set.add(originalIndividualStrings[i])) {
                throw new IllegalArgumentException(
                    "Exception during decoding - repeated representation: " + originalIndividualStrings[i]);
            }
        }
    }

    private String getStringRepresentationAndAddPaddingCharsIfNecessary(int index) {

        String paddingAndIndex = TWO_PADDING_CHARACTERS + index;
        return paddingAndIndex.substring(paddingAndIndex.length() - NUM_CHARS_PER_ENCODED_CHARACTER);
    }

    public BigInteger squareToEnsureIsGroupElement(BigInteger value, BigInteger p, BigInteger q) {

        if (value.compareTo(q) > 0) {
            throw new IllegalArgumentException(
                "Encoded writein too big for the mathematical group. Value: " + value + ", q: " + q);
        }

        return value.pow(2).mod(p);
    }

    private int calculateOptimalInitialStringBuilderLength(List<BigInteger> encoding) {

        int entireEncodingAsSingleStringLength =
            String.join("", encoding.stream().map(BigInteger::toString).collect(Collectors.toList())).length();

        int encodingAfterSquareRootStringLength = (int)Math.ceil(entireEncodingAsSingleStringLength / 2.0);
        
        int possiblePaddingCharactersAtStartOfEachElement = encoding.size() * 2;

        return encodingAfterSquareRootStringLength + possiblePaddingCharactersAtStartOfEachElement;
    }

    private int determineNumPaddingCharsNeeded(String originalIntegerValue) {

        int numPaddingCharsNeeded =
            NUM_CHARS_PER_ENCODED_CHARACTER - (originalIntegerValue.length() % NUM_CHARS_PER_ENCODED_CHARACTER);
        if (numPaddingCharsNeeded == NUM_CHARS_PER_ENCODED_CHARACTER) {
            numPaddingCharsNeeded = 0;
        }

        return numPaddingCharsNeeded;
    }

    private String decode(String toBeDecoded) {

        int length = toBeDecoded.length();
        StringBuilder sb = new StringBuilder(length / NUM_CHARS_PER_ENCODED_CHARACTER);

        for (int i = 0; i < length; i = i + NUM_CHARS_PER_ENCODED_CHARACTER) {

            String substring = toBeDecoded.substring(i, i + NUM_CHARS_PER_ENCODED_CHARACTER);
            int intValue = Integer.parseInt(substring);
            char character = alphabet.charAt((intValue - 1) - OFFSET);
            sb.append(character);
        }
        return sb.toString();
    }

    private void performBasicValidationsOnEncodeInputs(List<String> fields, BigInteger p, BigInteger q) {

        for (String writein : fields) {
            String[] parts = writein.split(separator);
            if (parts.length != 2) {
                throw new IllegalArgumentException(
                    "There should be exactly one separator character in each writein string, but this was not the case for: "
                        + writein);
            }
        }

        if (p == null) {
            throw new IllegalArgumentException("The p parameter was null");
        }
        if (q == null) {
            throw new IllegalArgumentException("The q parameter was null");
        }
    }

    private void performBasicValidationsOnDecodeInputs(BigInteger p) {

        if (p == null) {
            throw new IllegalArgumentException("The p parameter was null");
        }
    }
}
