package info.FrostFS.sdk;

import com.google.protobuf.AbstractMessage;
import com.google.protobuf.ByteString;
import frostFS.session.Types;
import info.FrostFS.sdk.tools.ECDsa;
import org.apache.commons.codec.digest.DigestUtils;
import org.bouncycastle.asn1.sec.SECNamedCurves;
import org.bouncycastle.asn1.sec.SECObjectIdentifiers;
import org.bouncycastle.crypto.params.ECDomainParameters;
import org.bouncycastle.crypto.params.ECPrivateKeyParameters;
import org.bouncycastle.crypto.signers.ECDSASigner;
import org.bouncycastle.crypto.signers.HMacDSAKCalculator;

import java.math.BigInteger;
import java.security.Signature;

import static org.bouncycastle.crypto.util.DigestFactory.createSHA256;
import static org.bouncycastle.util.BigIntegers.asUnsignedByteArray;

public class RequestSigner {
    public static final int RFC6979_SIGNATURE_SIZE = 64;

    public static byte[] signData(ECDsa key, byte[] data) {
        var hash = new byte[65];
        hash[0] = 0x04;
        try {
            Signature signature = Signature.getInstance("NONEwithECDSAinP1363Format");
            signature.initSign(key.getPrivateKey());
            signature.update(DigestUtils.sha512(data));
            byte[] sig = signature.sign();
            System.arraycopy(sig, 0, hash, 1, sig.length);
        } catch (Exception exp) {
            throw new RuntimeException(exp);
        }

        return hash;
    }

    public static byte[] signRFC6979(ECDsa key, byte[] data) {
        var digest = createSHA256();
        var secp256R1 = SECNamedCurves.getByOID(SECObjectIdentifiers.secp256r1);

        var ecParameters = new ECDomainParameters(secp256R1.getCurve(), secp256R1.getG(), secp256R1.getN());
        var ecPrivateKey = new ECPrivateKeyParameters(new BigInteger(1, key.getPrivateKeyByte()), ecParameters);
        var signer = new ECDSASigner(new HMacDSAKCalculator(digest));
        var hash = new byte[digest.getDigestSize()];

        digest.update(data, 0, data.length);
        digest.doFinal(hash, 0);
        signer.init(true, ecPrivateKey);

        var rs = signer.generateSignature(hash);
        var rBytes = asUnsignedByteArray(rs[0]);
        var sBytes = asUnsignedByteArray(rs[1]);

        var signature = new byte[RFC6979_SIGNATURE_SIZE];
        var index = RFC6979_SIGNATURE_SIZE / 2 - rBytes.length;
        System.arraycopy(rBytes, 0, signature, index, rBytes.length);
        index = RFC6979_SIGNATURE_SIZE - sBytes.length;
        System.arraycopy(sBytes, 0, signature, index, sBytes.length);
        return signature;
    }

    public static frostFS.refs.Types.SignatureRFC6979 signRFC6979(ECDsa key, AbstractMessage message) {
        return frostFS.refs.Types.SignatureRFC6979.newBuilder()
                .setKey(ByteString.copyFrom(key.getPublicKeyByte()))
                .setSign(ByteString.copyFrom(signRFC6979(key, message.toByteArray())))
                .build();
    }

    public static frostFS.refs.Types.SignatureRFC6979 signRFC6979(ECDsa key, ByteString data) {
        return frostFS.refs.Types.SignatureRFC6979.newBuilder()
                .setKey(ByteString.copyFrom(key.getPublicKeyByte()))
                .setSign(ByteString.copyFrom(signRFC6979(key, data.toByteArray())))
                .build();
    }

    public static frostFS.refs.Types.Signature signMessagePart(ECDsa key, AbstractMessage data) {
        var data2Sign = data.getSerializedSize() == 0 ? new byte[]{} : data.toByteArray();
        return frostFS.refs.Types.Signature.newBuilder()
                .setKey(ByteString.copyFrom(key.getPublicKeyByte()))
                .setSign(ByteString.copyFrom(signData(key, data2Sign)))
                .build();
    }

    public static void sign(AbstractMessage.Builder request, ECDsa key) {
        var meta = (AbstractMessage) request.getField(request.getDescriptorForType().findFieldByName("meta_header"));
        var body = (AbstractMessage) request.getField(request.getDescriptorForType().findFieldByName("body"));
        var verify = (AbstractMessage) request.getField(request.getDescriptorForType().findFieldByName("verify_header"));

        AbstractMessage.Builder verifyBuilder;
        if (verify instanceof Types.RequestVerificationHeader) {
            verifyBuilder = Types.RequestVerificationHeader.newBuilder();
        } else if (verify instanceof Types.ResponseVerificationHeader) {
            verifyBuilder = Types.ResponseVerificationHeader.newBuilder();
        } else {
            throw new IllegalArgumentException("Unsopported message type");
        }

        var verifyOrigin = (AbstractMessage) verify.getField(verify.getDescriptorForType().findFieldByName("origin"));
        if (verifyOrigin.getSerializedSize() == 0) {
            verifyBuilder.setField(
                    verifyBuilder.getDescriptorForType().findFieldByName("body_signature"),
                    signMessagePart(key, body)
            );
        } else {
            verifyBuilder.setField(verifyBuilder.getDescriptorForType().findFieldByName("origin"), verifyOrigin);
        }

        verifyBuilder.setField(
                verifyBuilder.getDescriptorForType().findFieldByName("meta_signature"),
                signMessagePart(key, meta)
        );
        verifyBuilder.setField(
                verifyBuilder.getDescriptorForType().findFieldByName("origin_signature"),
                signMessagePart(key, verifyOrigin)
        );

        request.setField(request.getDescriptorForType().findFieldByName("verify_header"), verifyBuilder.build());
    }
}