secp256k1-kmp/tests/src/commonTest/kotlin/fr/acinq/secp256k1/FrostTest.kt

package fr.acinq.secp256k1

import kotlinx.serialization.json.*

import kotlin.test.*

class FrostTest: BaseTest() {
    val msg32 = "this_could_be_the_hash_of_a_msg!".encodeToByteArray()

    @Test
    fun `frost share generation test cases`() {
        val tests = readData("frost/share_gen_vectors.json")
        val pubkeys = tests.jsonObject["pubkeys"]!!.jsonArray.map { Hex.decode(it.jsonPrimitive.content) }

        tests.jsonObject["valid_share_gen_test_cases"]!!.jsonArray.forEach { validTestCases ->

            val keyIndices = validTestCases.jsonObject["key_indices"]!!.jsonArray.map { it.jsonPrimitive.int }

            val seed32 = Hex.decode(validTestCases.jsonObject["seed"]!!.jsonPrimitive.content)
            val nParticipants = keyIndices.size
            val threshold = validTestCases.jsonObject["threshold"]!!.jsonPrimitive.int
            val ids33 = keyIndices.map { pubkeys[it] }.toTypedArray()

            val result = Secp256k1.frostSharesGen(
                seed32,
                threshold,
                nParticipants,
                ids33
            )

            val expected = validTestCases.jsonObject["expected"]!!;

            val expectedShare = expected.jsonObject["frost_share"]!!.jsonArray.map { Hex.decode(it.jsonPrimitive.content) }
            val expectedVSSCommitment = expected.jsonObject["vss_commitment"]!!.jsonArray.map { Hex.decode(it.jsonPrimitive.content) }
            val expectedPoK64 = Hex.decode(expected.jsonObject["pok64"]!!.jsonPrimitive.content)


            result.first.forEachIndexed { index, share ->
                assertEquals(
                    expected = Hex.encode(expectedShare[index]),
                    actual = Hex.encode(share),
                    "Unexpected $index:share for $keyIndices test case"
                )
            }
            result.second.forEachIndexed { index, vssCommitment ->
                assertEquals(
                    expected = Hex.encode(expectedVSSCommitment[index]),
                    actual = Hex.encode(vssCommitment),
                    "Unexpected $index:vss_commitment for the $keyIndices test case"
                )
            }
            assertEquals(
                expected = Hex.encode(expectedPoK64),
                actual = Hex.encode(result.third),
                message = "Unexpected pok64 for $keyIndices test case"
            )
        }
    }

    @Test
    fun `frost share generation signers`() {
        val tests = readData("frost/share_gen_signers_vectors.json")
        val pubkeys = tests.jsonObject["pubkeys"]!!.jsonArray.map { Hex.decode(it.jsonPrimitive.content) }

        val signerShareGenTestCase = tests.jsonObject["valid_signers_share_gen_test_case"]!!

        val keyIndices = signerShareGenTestCase.jsonObject["key_indices"]!!.jsonArray.map { it.jsonPrimitive.int }

        val nParticipants = keyIndices.size
        val threshold = signerShareGenTestCase.jsonObject["threshold"]!!.jsonPrimitive.int
        val ids33 = keyIndices.map { pubkeys[it] }.toTypedArray()

        signerShareGenTestCase.jsonObject["signers"]!!.jsonArray.forEachIndexed { signerIndex, signer ->
            val seed32 = Hex.decode(signer.jsonObject["seed"]!!.jsonPrimitive.content)
            println("Testing ${signer.jsonObject["seed"]!!.jsonPrimitive.content}")
            // There seems to be a bug that causes a crash if we call frost_share_gen too often

            val result = Secp256k1.frostSharesGen(
                seed32,
                threshold,
                nParticipants,
                ids33
            )

            val expected = signer.jsonObject["expected"]!!;

            val expectedShare = expected.jsonObject["frost_share"]!!.jsonArray.map { Hex.decode(it.jsonPrimitive.content) }
            val expectedVSSCommitment = expected.jsonObject["vss_commitment"]!!.jsonArray.map { Hex.decode(it.jsonPrimitive.content) }
            val expectedPoK64 = Hex.decode(expected.jsonObject["pok64"]!!.jsonPrimitive.content)

            result.first.forEachIndexed { index, share ->
                assertEquals(
                    expected = Hex.encode(expectedShare[index]),
                    actual = Hex.encode(share),
                    "Unexpected $signerIndex:signer $index:share for $keyIndices test case"
                )
            }
            result.second.forEachIndexed { index, vssCommitment ->
                assertEquals(
                    expected = Hex.encode(expectedVSSCommitment[index]),
                    actual = Hex.encode(vssCommitment),
                    "Unexpected $signerIndex:signer $index:vss_commitment for the $keyIndices test case"
                )
            }
            assertEquals(
                expected = Hex.encode(expectedPoK64),
                actual = Hex.encode(result.third),
                message = "Unexpected $signerIndex:signer pok64 for $keyIndices test case"
            )
        }
    }

    @Test
    fun `frost share aggregation`() {
        val shareGenTests = readData("frost/share_gen_vectors.json")
        val tests = readData("frost/share_agg_vectors.json")

        val expectedAggregatePublicKey = tests.jsonObject["aggregate_public_key"]!!.jsonPrimitive.content
        val publicKeys = shareGenTests.jsonObject["pubkeys"]!!.jsonArray.map { Hex.decode(it.jsonPrimitive.content) }

        val signerShareGenTestCase = shareGenTests.jsonObject["valid_signers_share_gen_test_case"]!!;

        val keyIndices = signerShareGenTestCase.jsonObject["key_indices"]!!.jsonArray.map { it.jsonPrimitive.int }
        val nParticipants = keyIndices.size
        val threshold = signerShareGenTestCase.jsonObject["threshold"]!!.jsonPrimitive.int
        val ids33 = keyIndices.map { publicKeys[it] }.toTypedArray()

        val vssCommitments = signerShareGenTestCase.jsonObject["signers"]!!.jsonArray.map { signer ->
            signer.jsonObject["expected"]!!.jsonObject["vss_commitment"]!!.jsonArray.map {
                Hex.decode(it.jsonPrimitive.content)
            }.toTypedArray()
        }

        signerShareGenTestCase.jsonObject["signers"]!!.jsonArray.forEachIndexed { index, _ ->
            val assignedShares = signerShareGenTestCase.jsonObject["signers"]!!.jsonArray.map {
                Hex.decode(
                    it.jsonObject["expected"]!!.jsonObject["frost_share"]!!.jsonArray[index].jsonPrimitive.content
                )
            }

            val result = Secp256k1.frostShareAggregate(
                assignedShares.toTypedArray(),
                vssCommitments.toTypedArray(),
                nParticipants,
                threshold,
                ids33[index]
            )

            val expected = tests.jsonObject["expected"]!!.jsonArray[index];

            val expectedAggregateShare = expected.jsonObject["aggregate_share"]!!.jsonPrimitive.content
            val expectedPublicShare = expected.jsonObject["public_share"]!!.jsonPrimitive.content

            assertEquals(
                expected = expectedAggregateShare,
                actual = Hex.encode(result.first),
                "Unexpected $index:aggregate_share"
            )
            assertEquals(
                expected = expectedAggregatePublicKey,
                actual = Hex.encode(result.second),
                "Unexpected $index:aggregate_public_key"
            )

            assertEquals(
                expected = 1,
                actual = Secp256k1.frostShareVerify(
                    threshold,
                    ids33[index],
                    assignedShares[index],
                    vssCommitments[index]
                ),
                message = "Couldn't verify share from $index signer"
            )

            assertEquals(
                expected = expectedPublicShare,
                actual = Hex.encode(
                    Secp256k1.frostComputePublicShare(
                        threshold,
                        ids33[index],
                        vssCommitments.toTypedArray(),
                        nParticipants
                    )
                ),
                message = "Couldn't verify share from $index signer"
            )
        }

    }

    @Test
    fun `frost share verify`() {
        val shareGenTests = readData("frost/share_gen_vectors.json")
        val tests = readData("frost/share_agg_vectors.json")

        val expectedAggregatePublicKey = tests.jsonObject["aggregate_public_key"]!!.jsonPrimitive.content
        val publicKeys = shareGenTests.jsonObject["pubkeys"]!!.jsonArray.map { Hex.decode(it.jsonPrimitive.content) }

        val signerShareGenTestCase = shareGenTests.jsonObject["valid_signers_share_gen_test_case"]!!;

        val keyIndices = signerShareGenTestCase.jsonObject["key_indices"]!!.jsonArray.map { it.jsonPrimitive.int }
        val nParticipants = keyIndices.size
        val threshold = signerShareGenTestCase.jsonObject["threshold"]!!.jsonPrimitive.int
        val ids33 = keyIndices.map { publicKeys[it] }.toTypedArray()

        val vssCommitments = signerShareGenTestCase.jsonObject["signers"]!!.jsonArray.map { signer ->
            signer.jsonObject["expected"]!!.jsonObject["vss_commitment"]!!.jsonArray.map {
                Hex.decode(it.jsonPrimitive.content)
            }.toTypedArray()
        }

        signerShareGenTestCase.jsonObject["signers"]!!.jsonArray.forEachIndexed { index, _ ->
            val assignedShares = signerShareGenTestCase.jsonObject["signers"]!!.jsonArray.map {
                Hex.decode(
                    it.jsonObject["expected"]!!.jsonObject["frost_share"]!!.jsonArray[index].jsonPrimitive.content
                )
            }

            val result = Secp256k1.frostShareAggregate(
                assignedShares.toTypedArray(),
                vssCommitments.toTypedArray(),
                nParticipants,
                threshold,
                ids33[index]
            )

            val expected = tests.jsonObject["expected"]!!.jsonArray[index];

            val expectedAggregateShare = expected.jsonObject["aggregate_share"]!!.jsonPrimitive.content
            val expectedPublicShare = expected.jsonObject["public_share"]!!.jsonPrimitive.content

            assertEquals(
                expected = expectedAggregateShare,
                actual = Hex.encode(result.first),
                "Unexpected $index:aggregate_share"
            )
            assertEquals(
                expected = expectedAggregatePublicKey,
                actual = Hex.encode(result.second),
                "Unexpected $index:aggregate_public_key"
            )

            assertEquals(
                expected = 1,
                actual = Secp256k1.frostShareVerify(
                    threshold,
                    ids33[index],
                    assignedShares[index],
                    vssCommitments[index]
                ),
                message = "Couldn't verify share from $index signer"
            )

            assertEquals(
                expected = expectedPublicShare,
                actual = Hex.encode(
                    Secp256k1.frostComputePublicShare(
                        threshold,
                        ids33[index],
                        vssCommitments.toTypedArray(),
                        nParticipants
                    )
                ),
                message = "Couldn't verify share from $index signer"
            )
        }
    }

    @Test
    fun `frost compute pubshare`() {
        val shareGenTests = readData("frost/share_gen_vectors.json")
        val tests = readData("frost/share_agg_vectors.json")

        val publicKeys = shareGenTests.jsonObject["pubkeys"]!!.jsonArray.map { Hex.decode(it.jsonPrimitive.content) }

        val signerShareGenTestCase = shareGenTests.jsonObject["valid_signers_share_gen_test_case"]!!;

        val keyIndices = signerShareGenTestCase.jsonObject["key_indices"]!!.jsonArray.map { it.jsonPrimitive.int }
        val nParticipants = keyIndices.size
        val threshold = signerShareGenTestCase.jsonObject["threshold"]!!.jsonPrimitive.int
        val ids33 = keyIndices.map { publicKeys[it] }.toTypedArray()

        val vssCommitments = signerShareGenTestCase.jsonObject["signers"]!!.jsonArray.map { signer ->
            signer.jsonObject["expected"]!!.jsonObject["vss_commitment"]!!.jsonArray.map {
                Hex.decode(it.jsonPrimitive.content)
            }.toTypedArray()
        }

        signerShareGenTestCase.jsonObject["signers"]!!.jsonArray.forEachIndexed { index, _ ->

            val expected = tests.jsonObject["expected"]!!.jsonArray[index];

            val expectedPublicShare = expected.jsonObject["public_share"]!!.jsonPrimitive.content

            assertEquals(
                expected = expectedPublicShare,
                actual = Hex.encode(
                    Secp256k1.frostComputePublicShare(
                        threshold,
                        ids33[index],
                        vssCommitments.toTypedArray(),
                        nParticipants
                    )
                ),
                message = "Couldn't verify share from $index signer"
            )
        }
    }

    // Frost Tweak
    @Test
    fun `frost pubkey tweak`() {
        val tests = readData("frost/share_agg_vectors.json")

        val expectedAggregatePublicKey = Hex.decode(
            tests.jsonObject["aggregate_public_key"]!!.jsonPrimitive.content
        )

        val expectedTweakCache = tests.jsonObject["tweak_cache"]!!.jsonPrimitive.content

        assertEquals(
            expected = expectedTweakCache,
            actual = Hex.encode(
                Secp256k1.frostPublicKeyTweak(
                    expectedAggregatePublicKey
                )
            ),
            message = "Tweak cache incorrect."
        )
    }

    @Test
    fun `frost pubkey ec tweak add`() {
        val tests = readData("frost/frost_tweak_vectors.json")

        val tweakCache = Hex.decode(
            tests.jsonObject["tweak_cache"]!!.jsonPrimitive.content
        )
        val ordinaryTweak = "this could be a BIP32 tweak.....".encodeToByteArray()

        assertEquals(
            expected = tests.jsonObject["ec_tweak_add"]!!.jsonPrimitive.content,
            actual = Secp256k1.frostPublicKeyEcTweakAdd(
                tweakCache,
                ordinaryTweak
            )?.let {
                Hex.encode(
                    it
                )
            },
            message = "EC Tweak Add incorrect."
        )
    }

    @Test
    fun `frost pubkey xonly tweak add`() {
        val tests = readData("frost/frost_tweak_vectors.json")

        val tweakCache = Hex.decode(
            tests.jsonObject["tweak_cache"]!!.jsonPrimitive.content
        )
        val ordinaryTweak = "this could be a BIP32 tweak.....".encodeToByteArray()

        assertEquals(
            expected = tests.jsonObject["ec_tweak_add"]!!.jsonPrimitive.content,
            // TODO: Return public key
            actual = Secp256k1.frostPublicKeyXonlyTweakAdd(
                tweakCache,
                ordinaryTweak
            )?.let {
                Hex.encode(
                    it
                )
            },
            message = "EC Tweak Add incorrect."
        )
    }

    // Frost Sign functionality
    @Test
    fun `frost nonce gen`() {
        val tests = readData("frost/frost_nonce_vectors.json")

        val sessionId = Hex.decode(
            tests.jsonObject["session_id"]!!.jsonPrimitive.content
        )
        val aggregatePublicKey = Hex.decode(
            tests.jsonObject["aggregate_public_key"]!!.jsonPrimitive.content
        )

        tests.jsonObject["signers"]!!.jsonArray.forEachIndexed { index, signers ->
            val aggregateShare = Hex.decode(
                signers.jsonObject["aggregate_share"]!!.jsonPrimitive.content
            )

            val (secNonce, pubNonce) = Secp256k1.frostNonceGen(
                sessionId,
                aggregateShare,
                msg32,
                aggregatePublicKey,
                null
            )

            val expectedSecNonce = signers.jsonObject["secnonce"]!!.jsonPrimitive.content
            val expectedPubNonce = signers.jsonObject["pubnonce"]!!.jsonPrimitive.content

            assertEquals(
                expected = expectedSecNonce,
                actual = Hex.encode(secNonce),
                message = "Invalid $index:secnonce"
            )

            assertEquals(
                expected = expectedPubNonce,
                actual = Hex.encode(pubNonce),
                message = "Invalid $index:pubnonce"
            )
        }

    }

    @Test
    fun `frost nonce process`() {
        val tests = readData("frost/frost_nonce_vectors.json")

        val pubnonces = tests.jsonObject["signers"]!!.jsonArray.map {
            Hex.decode(it.jsonObject["pubnonce"]!!.jsonPrimitive.content)
        }

        val aggregatePublicKey = Hex.decode(
            tests.jsonObject["aggregate_public_key"]!!.jsonPrimitive.content
        )

        val pubkeys = tests.jsonObject["pubkeys"]!!.jsonArray.map { Hex.decode(it.jsonPrimitive.content) }

        val tweakCache = Hex.decode(
            tests.jsonObject["tweak_cache"]!!.jsonPrimitive.content
        )

        tests.jsonObject["signers"]!!.jsonArray.forEachIndexed { signerIndex, signer ->

            val session = Secp256k1.frostNonceProcess(
                pubnonces.toTypedArray(),
                msg32,
                aggregatePublicKey,
                pubkeys[signerIndex],
                pubkeys.toTypedArray(),
                tweakCache,
                null
            )

            assertEquals(
                expected = signer.jsonObject["session"]!!.jsonPrimitive.content,
                actual = Hex.encode(session),
                message = "Invalid $signerIndex:session"
            )
        }
    }

    @Test
    fun `frost partial sign`() {
        val tests = readData("frost/frost_nonce_vectors.json")


        val tweakCache = Hex.decode(
            tests.jsonObject["tweak_cache"]!!.jsonPrimitive.content
        )

        tests.jsonObject["signers"]!!.jsonArray.forEachIndexed { signerIndex, signer ->

            val secNonce = Hex.decode(
                signer.jsonObject["secnonce"]!!.jsonPrimitive.content
            )

            val aggregateShare = Hex.decode(
                signer.jsonObject["aggregate_share"]!!.jsonPrimitive.content
            )

            val session = Hex.decode(
                signer.jsonObject["session"]!!.jsonPrimitive.content
            )


            val partialSignature = Secp256k1.frostPartialSign(
                secNonce,
                aggregateShare,
                session,
                tweakCache
            )

            assertEquals(
                expected = signer.jsonObject["partial_signature"]!!.jsonPrimitive.content,
                actual = Hex.encode(partialSignature),
                message = "Invalid $signerIndex:partial signature"
            )
        }
    }

    @Test
    fun `frost partial signature verify`() {
        val tests = readData("frost/frost_nonce_vectors.json")

        val tweakCache = Hex.decode(
            tests.jsonObject["tweak_cache"]!!.jsonPrimitive.content
        )

        tests.jsonObject["signers"]!!.jsonArray.forEach { signer ->
            val partialSignature = Hex.decode(
                signer.jsonObject["partial_signature"]!!.jsonPrimitive.content
            )

            val pubNonce = Hex.decode(
                signer.jsonObject["pubnonce"]!!.jsonPrimitive.content
            )

            val publicShare = Hex.decode(
                signer.jsonObject["public_share"]!!.jsonPrimitive.content
            )

            val session = Hex.decode(
                signer.jsonObject["session"]!!.jsonPrimitive.content
            )

            assertEquals(
                expected = 1,
                actual = Secp256k1.frostPartialSignatureVerify(
                    partialSignature,
                    pubNonce,
                    publicShare,
                    session,
                    tweakCache
                ),
                message = "Failed to verify partial signature"
            )
        }
    }

    @Test
    fun `frost partial signature aggregation`() {
        val tests = readData("frost/frost_nonce_vectors.json")

        val partialSignatures = tests.jsonObject["signers"]!!.jsonArray.map {
            Hex.decode(it.jsonObject["partial_signature"]!!.jsonPrimitive.content)
        }

        tests.jsonObject["signers"]!!.jsonArray.forEach { signer ->
            val session = Hex.decode(
                signer.jsonObject["session"]!!.jsonPrimitive.content
            )

            val aggregatedSignature = Secp256k1.frostPartialSignatureAggregate(
                session,
                partialSignatures.toTypedArray(),
            )

            assertEquals(
                expected = tests.jsonObject["aggregate_signature"]!!.jsonPrimitive.content,
                actual = Hex.encode(aggregatedSignature),
                message = "Invalid aggregate partial signature"
            )
        }
    }
}