#include <gtest/gtest.h>
extern "C" {
#include "../target/include/atms.h"
}

TEST(atms, produceAndVerifyAggregateSignature) {
    const char *msg = "some message";

    // Test with 5 parties and threshold 4.
    SigningKeyPtr sks[5];
    PublicKeyPtr keys[5];
    PublicKeyPoPPtr keys_pop[5];
    unsigned int nr_signers = 5;
    unsigned int threshold = 4;
    int err;

    for (int i = 0; i < nr_signers; i++) {
        err = atms_generate_keypair(&sks[i], &keys_pop[i]);
        atms_pkpop_to_pk(keys_pop[i], &keys[i]);
        ASSERT_EQ(err, 0);
    }

    // The threshold public key can be generated by using the public key parts of the participants.
    RegistrationPtr registration;
    AvkPtr avk_pk;

    err = avk_key_registration(keys_pop, nr_signers, &registration);
    ASSERT_EQ(err, 0);

    err = atms_registration_to_avk(&avk_pk, registration);
    ASSERT_EQ(err, 0);

    // Now, signers can produce threshold signatures with respect to `avk_pk`. Now, let's assume that
    // parties 0-3 generate a signature.
    SignaturePtr sigs[4];

    for (int i = 0; i < 4; i++) {
         err = atms_sign(msg, sks[i], &sigs[i]);
        ASSERT_EQ(err, 0);
    }

    // Given the four signatures and the public keys from the corresponding signers, any third party (not necessarily
    // a signer) can aggregate the signatures. It only needs knowledge of the avk key over which it is supposed to be
    // valid.
    AggregateSigPtr aggregated_sig;
    err = atms_aggregate_sigs(msg, sigs, keys, registration, 4, &aggregated_sig);
    ASSERT_EQ(err, 0);

    // Finally, we check that the signature is indeed valid.
    err = atms_verify_aggr(msg, aggregated_sig, avk_pk, threshold);
    ASSERT_EQ(err, 0);
}

TEST(atms, testingErrors) {
    const char *msg = "some message";

    // Test with 5 parties
    PublicKeyPoPPtr keys_pop[5];
    PublicKeyPtr keys[5];
    SigningKeyPtr sks[5];
    unsigned int nr_signers = 5;
    unsigned int threshold = 3;
    int err;

    for (int i = 0; i < nr_signers; i++) {
        err = atms_generate_keypair(&sks[i], &keys_pop[i]);
        atms_pkpop_to_pk(keys_pop[i], &keys[i]);
        ASSERT_EQ(err, 0);
    }

    // The threshold public key can be generated by using the public key parts of the participants. We leave out of this
    // the key of the "fake" party.
    RegistrationPtr registration;
    AvkPtr avk_pk;
    err = avk_key_registration(keys_pop, nr_signers, &registration);
    ASSERT_EQ(err, 0);

    err = atms_registration_to_avk(&avk_pk, registration);
    ASSERT_EQ(err, 0);

    // Now, signers can produce threshold signatures with respect to `avk_pk`. Now, let's assume that
    // only 0-3 generate a signature.
    SignaturePtr sigs[3];

    for (int i = 0; i < 3; i++) {
        err = atms_sign(msg, sks[i], &sigs[i]);
        ASSERT_EQ(err, 0);
    }

    // The signature is valid on its own
    err = atms_verify(msg, keys[0], sigs[0]);
    ASSERT_EQ(err, 0);

    // And the signature is invalid when verified over an invalid key
    err = atms_verify(msg, keys[1], sigs[0]);
    ASSERT_EQ(err, -1);

    // First we create a signature with no sufficient signers.
    AggregateSigPtr aggregated_sig_1;
    err = atms_aggregate_sigs(msg, sigs, keys, registration, 2, &aggregated_sig_1);
    ASSERT_EQ(err, 0);
    err = atms_verify_aggr(msg, aggregated_sig_1, avk_pk, threshold);
    ASSERT_EQ(err, -1);

    // We also create keys which will not be registered
    PublicKeyPoPPtr fake_keys_pop[3];
    PublicKeyPtr fake_keys[3];
    SigningKeyPtr fake_skeys[3];

    for (int i = 0; i < 3; i++) {
        err = atms_generate_keypair(&fake_skeys[i], &fake_keys_pop[i]);
        atms_pkpop_to_pk(fake_keys_pop[i], &fake_keys[i]);
        ASSERT_EQ(err, 0);
    }

    // We also create a "fake" avk;
    RegistrationPtr registration_fake;
    err = avk_key_registration(fake_keys_pop, 3, &registration_fake);
    ASSERT_EQ(err, 0);

    // Now, fake signature
    SignaturePtr fake_sigs[2];

    for (int i = 0; i < 2; i++) {
        err = atms_sign(msg, fake_skeys[i], &fake_sigs[i]);
        ASSERT_EQ(err, 0);
    }

    // Now we create a signature with fake_key and verify it with a different key. It will fail because the path
    // of the non-signers will be invalid with respect to the new key. Error message should be -3.
    AggregateSigPtr aggregated_sig_2;
    err = atms_aggregate_sigs(msg, fake_sigs, fake_keys, registration_fake, 2, &aggregated_sig_2);
    ASSERT_EQ(err, 0);
    err = atms_verify_aggr(msg, aggregated_sig_2, avk_pk, 3);
    ASSERT_EQ(err, -3);

    // If instead there are no `non-signers` (all signers submit a signature), but the signature is invalid, the error
    // message is -5, as the error now is simply that verification failed.
    PublicKeyPoPPtr fake_key_pop;
    PublicKeyPtr fake_key;
    SigningKeyPtr fake_skey;
    err = atms_generate_keypair(&fake_skey, &fake_key_pop);
    atms_pkpop_to_pk(fake_key_pop, &fake_key);
    ASSERT_EQ(err, 0);
    RegistrationPtr single_registration;
    err = avk_key_registration(&fake_key_pop, 1, &single_registration);
    ASSERT_EQ(err, 0);
    SignaturePtr fake_sig;
    err = atms_sign(msg, fake_skey, &fake_sig);
    ASSERT_EQ(err, 0);

    AggregateSigPtr aggregated_sig_3;
    err = atms_aggregate_sigs(msg, &fake_sig, &fake_key, single_registration, 1, &aggregated_sig_3);
    ASSERT_EQ(err, 0);
    err = atms_verify_aggr(msg, aggregated_sig_3, avk_pk, 3);
    ASSERT_EQ(err, -4);
}