import java.lang.foreign.Arena;
import java.lang.foreign.FunctionDescriptor;
import java.lang.foreign.Linker;
import java.lang.foreign.MemorySegment;
import java.lang.foreign.SymbolLookup;
import java.lang.invoke.MethodHandle;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.security.SecureRandom;
import static java.lang.foreign.ValueLayout.ADDRESS;
import static java.lang.foreign.ValueLayout.JAVA_INT;

/**
 * <p>compile with: -J-Duser.language=en
 * <p>run with: --enable-native-access=ALL-UNNAMED
 *
 * @see <a href="https://openjdk.org/jeps/434">FFI reference</a>
 */
@SuppressWarnings("OptionalGetWithoutIsPresent")
public class TestEcdsa {
	public static final int SECP256K1_FLAGS_TYPE_CONTEXT = 1;
	public static final int SECP256K1_CONTEXT_NONE = SECP256K1_FLAGS_TYPE_CONTEXT;
	public static final MethodHandle mh_secp256k1_context_create;
	public static final MethodHandle mh_secp256k1_context_randomize;
	public static final MethodHandle mh_secp256k1_context_destroy;
	public static final MethodHandle mh_secp256k1_ec_seckey_verify; // 验证随机数据是否可以作为私钥
	public static final MethodHandle mh_secp256k1_ec_pubkey_create; // 用私钥创建公钥
	public static final MethodHandle mh_secp256k1_ecdsa_sign_recoverable; // 用私钥为hash256生成签名
	public static final MethodHandle mh_secp256k1_ecdsa_recover; // 用签名和hash256还原公钥

	static {
		// https://github.com/bitcoin-core/secp256k1
		// x86_64-w64-mingw32-gcc.exe -DENABLE_MODULE_ECDH=1 -DENABLE_MODULE_RECOVERY=1 -DENABLE_MODULE_SCHNORRSIG=1 -DENABLE_MODULE_EXTRAKEYS=1 -DDLL_EXPORT -DUSE_FORCE_WIDEMUL_INT128=1 -DUSE_ASM_X86_64=1 -m64 -O3 -shared -o libsecp256k1.dll src/secp256k1.c src/precomputed_ecmult.c src/precomputed_ecmult_gen.c
		// cl /O2 /LD /MT -I. -DENABLE_MODULE_ECDH=1 -DENABLE_MODULE_RECOVERY=1 -DENABLE_MODULE_SCHNORRSIG=1 -DENABLE_MODULE_EXTRAKEYS=1 -DDLL_EXPORT src/secp256k1.c src/precomputed_ecmult.c src/precomputed_ecmult_gen.c
		var libSecp256k1 = SymbolLookup.libraryLookup("libsecp256k1.dll", Arena.global());
		var linker = Linker.nativeLinker();

		// secp256k1_context*(unsigned int flags)
		mh_secp256k1_context_create = linker.downcallHandle(
				libSecp256k1.find("secp256k1_context_create").get(),
				FunctionDescriptor.of(ADDRESS, JAVA_INT));

		// int(secp256k1_context* ctx, const unsigned char* seed32)
		mh_secp256k1_context_randomize = linker.downcallHandle(
				libSecp256k1.find("secp256k1_context_randomize").get(),
				FunctionDescriptor.of(JAVA_INT, ADDRESS, ADDRESS));

		// void(secp256k1_context* ctx)
		mh_secp256k1_context_destroy = linker.downcallHandle(
				libSecp256k1.find("secp256k1_context_destroy").get(),
				FunctionDescriptor.ofVoid(ADDRESS));

		// int(const secp256k1_context* ctx, const unsigned char* seckey)
		mh_secp256k1_ec_seckey_verify = linker.downcallHandle(
				libSecp256k1.find("secp256k1_ec_seckey_verify").get(),
				FunctionDescriptor.of(JAVA_INT, ADDRESS, ADDRESS));

		// int(const secp256k1_context* ctx, secp256k1_pubkey* pubkey, const unsigned char* seckey)
		mh_secp256k1_ec_pubkey_create = linker.downcallHandle(
				libSecp256k1.find("secp256k1_ec_pubkey_create").get(),
				FunctionDescriptor.of(JAVA_INT, ADDRESS, ADDRESS, ADDRESS));

		// int(const secp256k1_context* ctx, secp256k1_ecdsa_recoverable_signature* signature,
		// const unsigned char* msghash32, const unsigned char* seckey,
		// secp256k1_nonce_function noncefp, const void* noncedata)
		mh_secp256k1_ecdsa_sign_recoverable = linker.downcallHandle(
				libSecp256k1.find("secp256k1_ecdsa_sign_recoverable").get(),
				FunctionDescriptor.of(JAVA_INT, ADDRESS, ADDRESS, ADDRESS, ADDRESS, ADDRESS, ADDRESS));

		// int(const secp256k1_context* ctx, secp256k1_pubkey* pubkey,
		// const secp256k1_ecdsa_recoverable_signature* signature, const unsigned char* msghash32)
		mh_secp256k1_ecdsa_recover = linker.downcallHandle(
				libSecp256k1.find("secp256k1_ecdsa_recover").get(),
				FunctionDescriptor.of(JAVA_INT, ADDRESS, ADDRESS, ADDRESS, ADDRESS));
	}

	public static int num2Hex(int n) {
		return n + '0' + (((9 - n) >> 31) & ('A' - '9' - 1));
	}

	public static String toString(ByteBuffer bb) {
		int len = bb.limit();
		if (len <= 0)
			return "";
		var str = new byte[len * 3 - 1];
		for (int i = 0, j = 0; i < len; i++) {
			if (i > 0)
				str[j++] = '-';
			int b = bb.get(i);
			str[j++] = (byte)num2Hex((b >> 4) & 0xf);
			str[j++] = (byte)num2Hex(b & 0xf);
		}
		return new String(str, StandardCharsets.ISO_8859_1);
	}

	private static final SecureRandom secRand = new SecureRandom();

	public static void rand(ByteBuffer bb) {
		byte[] buf = new byte[bb.limit()];
		secRand.nextBytes(buf);
		bb.put(0, buf);
	}

	public static void main(String[] args) throws Throwable {
		try (Arena arena = Arena.ofConfined()) {
			var priKey = arena.allocate(32);
			var pubKey = arena.allocate(64);
			var msgHash = arena.allocate(32);
			var sigData = arena.allocate(65);
			var pubKey2 = arena.allocate(64);
			var priKeyBB = priKey.asByteBuffer();
			var pubKeyBB = pubKey.asByteBuffer();
			var msgHashBB = msgHash.asByteBuffer();
			var sigDataBB = sigData.asByteBuffer();
			var pubKey2BB = pubKey2.asByteBuffer();

			var ctx = (MemorySegment)mh_secp256k1_context_create.invokeExact(SECP256K1_CONTEXT_NONE);

			rand(priKeyBB);
			int r = (int)mh_secp256k1_context_randomize.invokeExact(ctx, priKey);
			System.out.println("secp256k1_context_randomize: " + r);

			do {
				rand(priKeyBB);
				r = (int)mh_secp256k1_ec_seckey_verify.invokeExact(ctx, priKey);
				System.out.println("secp256k1_ec_seckey_verify: " + r);
			} while (r != 1);
			System.out.println("priKey : [" + priKey.byteSize() + "] " + toString(priKeyBB));

			r = (int)mh_secp256k1_ec_pubkey_create.invokeExact(ctx, pubKey, priKey);
			System.out.println("secp256k1_ec_pubkey_create: " + r);
			System.out.println("pubKey : [" + pubKey.byteSize() + "] " + toString(pubKeyBB));

			rand(msgHashBB);
			r = (int)mh_secp256k1_ecdsa_sign_recoverable.invokeExact(ctx, sigData, msgHash, priKey,
					MemorySegment.NULL, MemorySegment.NULL);
			System.out.println("secp256k1_ecdsa_sign_recoverable: " + r);
			System.out.println("sigData: [" + sigData.byteSize() + "] " + toString(sigDataBB));

			r = (int)mh_secp256k1_ecdsa_recover.invokeExact(ctx, pubKey2, sigData, msgHash);
			System.out.println("mh_secp256k1_ecdsa_recover: " + r);
			System.out.println("pubKey2: [" + pubKey2.byteSize() + "] " + toString(pubKey2BB));

			System.out.println("checkPubKey: " + pubKeyBB.equals(pubKey2BB));

			mh_secp256k1_context_destroy.invokeExact(ctx);
			System.out.println("TEST OK!");

			for (int j = 0; j < 5; j++) {
				var t = System.nanoTime();
				for (int i = 0; i < 10000; i++) {
					r = (int)mh_secp256k1_ecdsa_recover.invokeExact(ctx, pubKey2, sigData, msgHash);
					if (r != 1 || !pubKeyBB.equals(pubKey2BB))
						throw new AssertionError();
				}
				System.out.println((System.nanoTime() - t) / 1_000_000 + " ms");
			}
		}
	}
}