package com.jscape.util.keystore;

import com.jscape.util.ArrayTools;
import com.jscape.util.Pair;
import com.jscape.util.crypto.CryptoServicesProvider;
import com.jscape.util.io.DataCodec;
import com.jscape.util.io.primitive.IntByteArrayCodec;
import com.jscape.util.io.primitive.IntStringCodec;
import org.bouncycastle.asn1.x9.ECNamedCurveTable;
import org.bouncycastle.asn1.x9.X9ECParameters;
import org.bouncycastle.jcajce.provider.asymmetric.util.EC5Util;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.security.AlgorithmParameters;
import java.security.KeyFactory;
import java.security.Provider;
import java.security.PublicKey;
import java.security.interfaces.ECPublicKey;
import java.security.spec.ECGenParameterSpec;
import java.security.spec.ECParameterSpec;
import java.security.spec.ECPoint;
import java.security.spec.ECPublicKeySpec;
import java.util.Enumeration;
import java.util.Map;

import static com.jscape.util.CollectionTools.map;
import static com.jscape.util.ExceptionTools.ioException;

/**
 * @author Alex Usun
 */
public class SshEcPublicKeyCodec
        implements DataCodec<PublicKey> {

    private static final int NO_COMPRESSION_CODE = 0x04;
    private static final Map<String, String> REQUIRED_CURVES = map(Pair.pair("1.2.840.10045.3.1.7", "nistp256"),
                                                                   Pair.pair("1.3.132.0.34", "nistp384"),
                                                                   Pair.pair("1.3.132.0.35", "nistp521"));

    private final Provider provider;

    public static void writeValue(PublicKey value, OutputStream out)
            throws IOException {
        try {
            ECPublicKey key = (ECPublicKey) value;
            ECPoint w = key.getW();
            byte[] r = w.getAffineX().toByteArray();
            byte[] s = w.getAffineY().toByteArray();
            byte[] q = new byte[1 + r.length + s.length];
            q[0] = NO_COMPRESSION_CODE;
            System.arraycopy(r, 0, q, 1, r.length);
            System.arraycopy(s, 0, q, 1 + r.length, s.length);

            String identifier = identifierFor(key);

            IntStringCodec.writeBigEndianValue(String.format("ecdsa-sha2-%s", identifier), StandardCharsets.US_ASCII, out);
            IntStringCodec.writeBigEndianValue(identifier, StandardCharsets.US_ASCII, out);
            IntByteArrayCodec.writeBigEndianValue(q, out);
        } catch (Throwable e) {
            throw ioException(e);
        }
    }

    public static PublicKey readValue(InputStream in, Provider provider)
            throws IOException {
        try {
            String algorithm = IntStringCodec.readBigEndianValue(StandardCharsets.US_ASCII, in);
            assertSupported(algorithm);

            String identifier = IntStringCodec.readBigEndianValue(StandardCharsets.US_ASCII, in);
            byte[] q = IntByteArrayCodec.readBigEndianValue(in);

            int entryLength = (q.length - 1) / 2;
            byte[] r = ArrayTools.copy(q, 1, new byte[entryLength]);
            byte[] s = ArrayTools.copy(q, 1 + entryLength, new byte[entryLength]);
            AlgorithmParameters algorithmParameters = provider != null ? AlgorithmParameters.getInstance("EC", provider) : AlgorithmParameters.getInstance("EC");
            algorithmParameters.init(new ECGenParameterSpec(oidFor(identifier)));
            ECParameterSpec ecParameterSpec = algorithmParameters.getParameterSpec(ECParameterSpec.class);
            ECPoint w = new ECPoint(new BigInteger(1, r), new BigInteger(1, s));
            ECPublicKeySpec keySpec = new ECPublicKeySpec(w, ecParameterSpec);

            KeyFactory keyFactory = provider != null ? KeyFactory.getInstance("EC", (Provider) provider) : KeyFactory.getInstance("EC");

            return keyFactory.generatePublic(keySpec);
        } catch (Throwable e) {
            throw ioException(e);
        }
    }

    private static void assertSupported(String algorithm)
            throws IOException {
        if (!algorithm.startsWith("ecdsa-sha2-")) {
            throw new IOException("Unsupported algorithm: " + algorithm);
        }
    }

    public static SshEcPublicKeyCodec defaultCodec() {
        return new SshEcPublicKeyCodec(null);
    }

    public static SshEcPublicKeyCodec bcCodec() {
        return new SshEcPublicKeyCodec(CryptoServicesProvider.instance());
    }

    private static String oidFor(String identifier) {
        for (Map.Entry<String, String> entry : REQUIRED_CURVES.entrySet()) {
            if (entry.getValue().equals(identifier)) {
                return entry.getKey();
            }
        }

        return identifier;
    }

    public static String identifierFor(ECPublicKey key) {
        String identifier = identifierFor(EC5Util.convertSpec((key.getParams())));
        return REQUIRED_CURVES.getOrDefault(identifier, identifier);
    }

    private static String identifierFor(org.bouncycastle.jce.spec.ECParameterSpec spec) {
        for (Enumeration names = ECNamedCurveTable.getNames(); names.hasMoreElements();) {
            String name = (String) names.nextElement();
            X9ECParameters params = ECNamedCurveTable.getByName(name);

            if (params.getN().equals(spec.getN()) && params.getH().equals(spec.getH()) && params.getCurve().equals(spec.getCurve()) && params.getG().equals(spec.getG())) {
                return ECNamedCurveTable.getOID(name).toString();
            }
        }

        throw new IllegalArgumentException(String.format("Unsupported curve type: %s.", spec));
    }

    public SshEcPublicKeyCodec(Provider provider) {
        this.provider = provider;
    }

    public void write(PublicKey value, OutputStream out)
            throws IOException {
        writeValue(value, out);
    }

    public PublicKey read(InputStream in)
            throws IOException {
        return readValue(in, this.provider);
    }

}