package fr.acinq.secp256k1

import kotlinx.serialization.json.*

import kotlin.test.*

class FrostTest: BaseTest() {

    @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()

            println("Testing ${Hex.encode(seed32)}")

            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 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 expectedPublicKey = expected.jsonObject["aggregate_public_key"]!!.jsonPrimitive.content
            val expectedPublicShare = expected.jsonObject["public_share"]!!.jsonPrimitive.content

            assertEquals(
                expected = expectedAggregateShare,
                actual = Hex.encode(result.first),
                "Unexpected $index:aggregate_share"
            )
            assertEquals(
                expected = expectedPublicKey,
                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 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 expectedPublicKey = expected.jsonObject["aggregate_public_key"]!!.jsonPrimitive.content
            val expectedPublicShare = expected.jsonObject["public_share"]!!.jsonPrimitive.content

            assertEquals(
                expected = expectedAggregateShare,
                actual = Hex.encode(result.first),
                "Unexpected $index:aggregate_share"
            )
            assertEquals(
                expected = expectedPublicKey,
                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`() {

    }

    @Test
    fun `frost pubkey ec tweak add`() {

    }

    @Test
    fun `frost pubkey xonly tweak add`() {

    }

    // Frost Sign functionality
    @Test
    fun `frost nonce gen`() {

    }

    @Test
    fun `frost nonce process`() {

    }

    @Test
    fun `frost partial sign `() {

    }

    @Test
    fun `frost partial signature verify`() {

    }

    @Test
    fun `frost partial signature aggregation`() {

    }
}