package org.keycloak.protocol.docker;

import org.keycloak.models.utils.Base32;

import java.security.Key;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Set;
import java.util.function.BiConsumer;
import java.util.function.BinaryOperator;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collector;
import java.util.stream.Stream;

/**
 * The “kid” field has to be in a libtrust fingerprint compatible format. Such a format can be generated by following steps:
 * 1) Take the DER encoded public key which the JWT token was signed against.
 * 2) Create a SHA256 hash out of it and truncate to 240bits.
 * 3) Split the result into 12 base32 encoded groups with : as delimiter.
 *
 * Ex: "kid": "PYYO:TEWU:V7JH:26JV:AQTZ:LJC3:SXVJ:XGHA:34F2:2LAQ:ZRMK:Z7Q6"
 *
 * @see https://docs.docker.com/registry/spec/auth/jwt/
 * @see https://github.com/docker/libtrust/blob/master/key.go#L24
 */
public class DockerKeyIdentifier {

    private final String identifier;

    public DockerKeyIdentifier(final Key key) throws InstantiationException {
        try {
            final MessageDigest sha256 = MessageDigest.getInstance("SHA-256");
            final byte[] hashed = sha256.digest(key.getEncoded());
            final byte[] hashedTruncated = truncateToBitLength(240, hashed);
            final String base32Id = Base32.encode(hashedTruncated);
            identifier = byteStream(base32Id.getBytes()).collect(new DelimitingCollector());
        } catch (final NoSuchAlgorithmException e) {
            throw new InstantiationException("Could not instantiate docker key identifier, no SHA-256 algorithm available.");
        }
    }

    // ugh.
    private Stream<Byte> byteStream(final byte[] bytes) {
        final Collection<Byte> colectionedBytes = new ArrayList<>();
        for (final byte aByte : bytes) {
            colectionedBytes.add(aByte);
        }

        return colectionedBytes.stream();
    }

    private byte[] truncateToBitLength(final int bitLength, final byte[] arrayToTruncate) {
        if (bitLength % 8 != 0) {
            throw new IllegalArgumentException("Bit length for truncation of byte array given as a number not divisible by 8");
        }

        final int numberOfBytes = bitLength / 8;
        return Arrays.copyOfRange(arrayToTruncate, 0, numberOfBytes);
    }

    @Override
    public String toString() {
        return identifier;
    }

    @Override
    public boolean equals(final Object o) {
        if (this == o) return true;
        if (!(o instanceof DockerKeyIdentifier)) return false;

        final DockerKeyIdentifier that = (DockerKeyIdentifier) o;

        return identifier != null ? identifier.equals(that.identifier) : that.identifier == null;

    }

    @Override
    public int hashCode() {
        return identifier != null ? identifier.hashCode() : 0;
    }

    // Could probably be generalized with size and delimiter arguments, but leaving it here for now until someone else needs it.
    public static class DelimitingCollector implements Collector<Byte, StringBuilder, String> {

        @Override
        public Supplier<StringBuilder> supplier() {
            return () -> new StringBuilder();
        }

        @Override
        public BiConsumer<StringBuilder, Byte> accumulator() {
            return ((stringBuilder, aByte) -> {
                if (needsDelimiter(4, ":", stringBuilder)) {
                    stringBuilder.append(":");
                }

                stringBuilder.append(new String(new byte[]{aByte}));
            });
        }

        private static boolean needsDelimiter(final int maxLength, final String delimiter, final StringBuilder builder) {
            final int lastDelimiter = builder.lastIndexOf(delimiter);
            final int charsSinceLastDelimiter = builder.length() - lastDelimiter;
            return charsSinceLastDelimiter > maxLength;
        }

        @Override
        public BinaryOperator<StringBuilder> combiner() {
            return ((left, right) -> new StringBuilder(left.toString()).append(right.toString()));
        }

        @Override
        public Function<StringBuilder, String> finisher() {
            return StringBuilder::toString;
        }

        @Override
        public Set<Characteristics> characteristics() {
            return Collections.emptySet();
        }
    }
}
