package com.jscape.inet.ssh.protocol.v2.marshaling.keys;

import com.jscape.inet.ssh.protocol.v2.marshaling.primitive.StringCodec;
import com.jscape.inet.ssh.protocol.v2.messages.NamedEllipticCurveDomainParameters;
import com.jscape.util.ArrayTools;
import com.jscape.util.io.DataCodec;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.math.BigInteger;
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;


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

    private static final int NO_COMPRESSION_CODE = 0x04;

    private final Object provider;

    public static EcPublicKeyCodec codecFor(String provider) {
        return new EcPublicKeyCodec(provider);
    }

    public static EcPublicKeyCodec codecFor(Provider provider) {
        return new EcPublicKeyCodec(provider);
    }

    public static PublicKey readValue(String identifier, Object provider, InputStream in)
            throws IOException {
        byte[] q = StringCodec.readValue(in);
        assertQValid(q);

        try {
            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 instanceof String ? AlgorithmParameters.getInstance("EC", (String) provider) : provider instanceof Provider ? AlgorithmParameters.getInstance("EC", (Provider) provider) : AlgorithmParameters.getInstance("EC");
            algorithmParameters.init(new ECGenParameterSpec(NamedEllipticCurveDomainParameters.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 instanceof String ? KeyFactory.getInstance("EC", (String) provider) : provider instanceof Provider ? KeyFactory.getInstance("EC", (Provider) provider) : KeyFactory.getInstance("EC");

            return keyFactory.generatePublic(keySpec);
        } catch (Exception ex) {
            throw new IOException(ex);
        }
    }

    public static void writeValue(PublicKey value, OutputStream out)
            throws IOException {
        ECPublicKey publicKey = (ECPublicKey) value;
        int elementSize = (publicKey.getParams().getCurve().getField().getFieldSize() + 7) / 8;
        ECPoint w = publicKey.getW();
        byte[] r = removeLeadingZeroes(w.getAffineX().toByteArray());
        byte[] s = removeLeadingZeroes(w.getAffineY().toByteArray());
        byte[] q = new byte[1 + elementSize * 2];
        q[0] = NO_COMPRESSION_CODE;
        System.arraycopy(r, 0, q, 1 + elementSize - r.length, r.length);
        System.arraycopy(s, 0, q,  1 + elementSize + elementSize - s.length, s.length);

        StringCodec.writeValue(q, out);
    }

    private static void assertQValid(byte[] q)
            throws IOException {
        if (q[0] != NO_COMPRESSION_CODE) {
            throw new IOException("Compressed EC public key value is not supported.");
        }
    }

    private static byte[] removeLeadingZeroes(byte[] data) {
        int i = 0;
        while (data[i] == 0x00) {
            ++i;
        }

        return i == 0 ? data : ArrayTools.copy(data, i, new byte[data.length - i]);
    }

    private EcPublicKeyCodec(Object provider) {
        this.provider = provider;
    }

    @Override
    public PublicKey read(InputStream in)
            throws IOException {
        String identifier = StringCodec.readUsAsciiValue(in);
        return readValue(identifier, this.provider, in);
    }

    @Override
    public void write(PublicKey value, OutputStream out)
            throws IOException {
        StringCodec.writeUsAsciiValue(NamedEllipticCurveDomainParameters.identifierFor(value), out);
        writeValue(value, out);
    }
}
