/***********************************************************************
 * Distributed under the MIT software license, see the accompanying    *
 * file COPYING or https://www.opensource.org/licenses/mit-license.php.*
 ***********************************************************************/

#ifndef SECP256K1_MODULE_CHILLDKG_SIMPLPEDPOP_IMPL_H
#define SECP256K1_MODULE_CHILLDKG_SIMPLPEDPOP_IMPL_H

#include <string.h>

#include "../../../include/secp256k1.h"

#include "util.h"
#include "vss.h"
#include "simplpedpop.h"
#include "../../ecmult.h"
#include "../../util.h"

/* The tag prefix of the proof-of-possession signatures (POP_MSG_TAG in
 * simplpedpop.py). */
#define SECP256K1_CHILLDKG_POP_MSG_TAG "BIP DKG/pop message"

static int secp256k1_chilldkg_simplpedpop_pop_prove(const secp256k1_context *ctx, unsigned char *pop64, const unsigned char *seckey32, uint32_t participant_id, const unsigned char *aux_rand32) {
    /* pop_msg: the message is u32be(participant_id). */
    unsigned char msg[4];
    int ret;

    secp256k1_write_be32(msg, participant_id);
    ret = secp256k1_chilldkg_schnorrsig_sign(ctx, pop64, msg, sizeof(msg), seckey32, aux_rand32, SECP256K1_CHILLDKG_POP_MSG_TAG);
    secp256k1_memclear_explicit(msg, sizeof(msg));
    return ret;
}

static int secp256k1_chilldkg_simplpedpop_pop_verify(const secp256k1_context *ctx, const unsigned char *pop64, const unsigned char *pubkey32, uint32_t participant_id) {
    unsigned char msg[4];
    int ret;

    secp256k1_write_be32(msg, participant_id);
    ret = secp256k1_chilldkg_schnorrsig_verify(ctx, pop64, msg, sizeof(msg), pubkey32, SECP256K1_CHILLDKG_POP_MSG_TAG);
    secp256k1_memclear_explicit(msg, sizeof(msg));
    return ret;
}

static int secp256k1_chilldkg_simplpedpop_participant_step1(const secp256k1_context *ctx, secp256k1_chilldkg_simplpedpop_participant_state *state, unsigned char *msg, secp256k1_scalar *partial_secshares, const unsigned char *seed32, uint32_t t, uint32_t n, uint32_t participant_id, const unsigned char *aux_rand32) {
    const secp256k1_hash_ctx *hash_ctx;
    secp256k1_scalar coeffs[SECP256K1_CHILLDKG_MAX_PARTICIPANTS];
    secp256k1_ge coms[SECP256K1_CHILLDKG_MAX_PARTICIPANTS];
    unsigned char seckey32[32];
    unsigned char pop[64];
    uint32_t i;
    size_t j;
    int coeffs_ok;
    int ret = 0;

    VERIFY_CHECK(ctx != NULL);
    ARG_CHECK(secp256k1_ecmult_gen_context_is_built(&ctx->ecmult_gen_ctx));
    ARG_CHECK(state != NULL);
    ARG_CHECK(msg != NULL);
    ARG_CHECK(partial_secshares != NULL);
    ARG_CHECK(seed32 != NULL);
    ARG_CHECK(aux_rand32 != NULL);
    hash_ctx = secp256k1_get_hash_context(ctx);

    /* Input validation, mirroring the ValueError/IndexError checks of the
     * reference. The reference would fail on t == 0 when accessing the
     * constant coefficient; reject it here instead. */
    if (t < 1 || n < 1 || t > n || participant_id >= n || n > SECP256K1_CHILLDKG_MAX_PARTICIPANTS) {
        return 0;
    }

    /* VSS.generate(seed, t). */
    coeffs_ok = secp256k1_chilldkg_vss_gen_coeffs(hash_ctx, coeffs, t, seed32);
    /* The failure depends on the secret seed but occurs with negligible
     * probability only (the reference raises OverflowError in this case). */
    secp256k1_declassify(ctx, &coeffs_ok, sizeof(coeffs_ok));
    if (coeffs_ok) {
        /* vss.secshares(n): the partial secret shares from us to all
         * participants. */
        for (i = 0; i < n; i++) {
            secp256k1_chilldkg_vss_secshare_for(&partial_secshares[i], coeffs, t, i);
        }
        /* com = vss.commit(). */
        secp256k1_chilldkg_vss_commit(ctx, coms, coeffs, t);
        /* pop = pop_prove(vss.secret(), participant_id, aux_rand); the secret
         * is the constant coefficient f(0) = coeffs[0]. Fails only if the
         * secret is zero (negligible probability). */
        secp256k1_scalar_get_b32(seckey32, &coeffs[0]);
        if (secp256k1_chilldkg_simplpedpop_pop_prove(ctx, pop, seckey32, participant_id, aux_rand32)) {
            /* msg = ParticipantMsg(com, pop).to_bytes() */
            for (j = 0; j < t; j++) {
                secp256k1_chilldkg_point_save(msg + 33 * j, &coms[j]);
            }
            memcpy(msg + 33 * t, pop, 64);
            state->t = t;
            state->n = n;
            state->participant_id = participant_id;
            /* com_to_secret = com.commitment_to_secret() */
            state->com_to_secret = coms[0];
            ret = 1;
        }
    }

    for (j = 0; j < t; j++) {
        secp256k1_scalar_clear(&coeffs[j]);
    }
    secp256k1_memclear_explicit(seckey32, sizeof(seckey32));
    secp256k1_memclear_explicit(pop, sizeof(pop));
    return ret;
}

static void secp256k1_chilldkg_simplpedpop_participant_step2_prepare_secshare(secp256k1_scalar *out, const secp256k1_scalar *partial_secshares, size_t n) {
    size_t i;

    secp256k1_scalar_set_int(out, 0);
    for (i = 0; i < n; i++) {
        secp256k1_scalar_add(out, out, &partial_secshares[i]);
    }
}

static secp256k1_chilldkg_fault secp256k1_chilldkg_simplpedpop_participant_step2(const secp256k1_context *ctx, secp256k1_chilldkg_simplpedpop_dkg_output *dkg_output, unsigned char *eq_input, uint32_t *fault_index, secp256k1_chilldkg_simplpedpop_participant_inv_data *inv_data, const secp256k1_chilldkg_simplpedpop_participant_state *state, const unsigned char *cmsg, size_t cmsg_len, const secp256k1_scalar *secshare) {
    const uint32_t t = state->t;
    const uint32_t n = state->n;
    const uint32_t participant_id = state->participant_id;
    secp256k1_ge coms_to_secrets[SECP256K1_CHILLDKG_MAX_PARTICIPANTS];
    secp256k1_ge sum_coms_to_nonconst_terms[SECP256K1_CHILLDKG_MAX_PARTICIPANTS - 1];
    secp256k1_ge sum_coms[SECP256K1_CHILLDKG_MAX_PARTICIPANTS];
    secp256k1_ge sum_coms_tweaked[SECP256K1_CHILLDKG_MAX_PARTICIPANTS];
    secp256k1_ge pubtweak;
    secp256k1_ge pubshare;
    secp256k1_ge pubshare_i;
    secp256k1_gej pubsharej;
    secp256k1_scalar tweak;
    secp256k1_scalar secshare_tweaked;
    const unsigned char *pops;
    unsigned char buf33[33];
    unsigned char buf33b[33];
    unsigned char xonly32[32];
    uint32_t i;
    size_t j;

    VERIFY_CHECK(ctx != NULL);
    ARG_CHECK(secp256k1_ecmult_gen_context_is_built(&ctx->ecmult_gen_ctx));
    ARG_CHECK(dkg_output != NULL);
    ARG_CHECK(eq_input != NULL);
    ARG_CHECK(state != NULL);
    ARG_CHECK(cmsg != NULL);
    ARG_CHECK(secshare != NULL);

    /* CoordinatorMsg.from_bytes; parse errors blame the coordinator. */
    if (cmsg_len != 97 * (size_t)n + 33 * (size_t)(t - 1)) {
        return SECP256K1_CHILLDKG_FAULTY_COORDINATOR;
    }
    for (i = 0; i < n; i++) {
        if (!secp256k1_chilldkg_point_load(&coms_to_secrets[i], cmsg + 33 * i)) {
            return SECP256K1_CHILLDKG_FAULTY_COORDINATOR;
        }
    }
    /* Parse the summed commitments to the non-constant terms. */
    for (j = 1; j < t; j++) {
        if (!secp256k1_chilldkg_point_load(&sum_coms_to_nonconst_terms[j - 1], cmsg + 33 * n + 33 * (j - 1))) {
            return SECP256K1_CHILLDKG_FAULTY_COORDINATOR;
        }
    }
    pops = cmsg + 33 * n + 33 * (t - 1);

    /* The coordinator must echo our own commitment to the secret. Compare the
     * canonical serializations so that the point at infinity (which cannot
     * occur here in practice, see participant_step1) is handled, too. */
    secp256k1_chilldkg_point_save(buf33, &state->com_to_secret);
    secp256k1_chilldkg_point_save(buf33b, &coms_to_secrets[participant_id]);
    if (secp256k1_memcmp_var(buf33, buf33b, 33) != 0) {
        return SECP256K1_CHILLDKG_FAULTY_COORDINATOR;
    }

    for (i = 0; i < n; i++) {
        if (i == participant_id) {
            /* No need to check our own pop. */
            continue;
        }
        if (secp256k1_ge_is_infinity(&coms_to_secrets[i])) {
            /* "Participant sent invalid commitment": the x-only serialization
             * needed for the pop verification does not exist. */
            *fault_index = i;
            return SECP256K1_CHILLDKG_FAULTY_PARTICIPANT_OR_COORDINATOR;
        }
        secp256k1_chilldkg_xonly_save(xonly32, &coms_to_secrets[i]);
        if (!secp256k1_chilldkg_simplpedpop_pop_verify(ctx, pops + 64 * i, xonly32, i)) {
            /* "Participant sent invalid proof-of-knowledge". */
            *fault_index = i;
            return SECP256K1_CHILLDKG_FAULTY_PARTICIPANT_OR_COORDINATOR;
        }
    }

    secp256k1_chilldkg_simplpedpop_assemble_sum_coms(sum_coms, coms_to_secrets, sum_coms_to_nonconst_terms, n, t);

    if (!secp256k1_chilldkg_vss_invalid_taproot_commit(ctx, sum_coms_tweaked, &tweak, &pubtweak, sum_coms, t)) {
        /* Unreachable in practice: sum_coms[0] can only be the point at
         * infinity if some pop is invalid (checked above), and a tweak hash
         * overflow has negligible probability. The reference raises a
         * non-protocol error here. Fill inv_data with the untweaked values so
         * that the investigation procedure can still run. */
        *fault_index = UINT32_MAX;
        if (inv_data != NULL) {
            inv_data->n = n;
            inv_data->participant_id = participant_id;
            inv_data->secshare = *secshare;
            secp256k1_chilldkg_vss_pubshare(&pubsharej, sum_coms, t, participant_id);
            secp256k1_ge_set_gej_var(&inv_data->pubshare, &pubsharej);
        }
        return SECP256K1_CHILLDKG_UNKNOWN_FAULTY_PARTICIPANT_OR_COORDINATOR;
    }

    /* Verifying the tweaked secshare against the tweaked pubshare is
     * equivalent to verifying the untweaked values, but avoids computing the
     * untweaked pubshare in the happy path. */
    secp256k1_chilldkg_vss_pubshare(&pubsharej, sum_coms_tweaked, t, participant_id);
    secp256k1_ge_set_gej_var(&pubshare, &pubsharej);
    secp256k1_scalar_add(&secshare_tweaked, secshare, &tweak);
    if (!secp256k1_chilldkg_vss_verify_secshare(ctx, &secshare_tweaked, &pubshare)) {
        /* "Received invalid secshare"; capture the investigation data so
         * that the investigation procedure can determine the faulty party. */
        *fault_index = UINT32_MAX;
        if (inv_data != NULL) {
            secp256k1_ge pubtweak_neg;
            inv_data->n = n;
            inv_data->participant_id = participant_id;
            /* The investigation works on the untweaked values. */
            inv_data->secshare = *secshare;
            /* pubshare = pubshare_tweaked - pubtweak (gej_add_ge_var handles
             * a pubtweak at infinity, which occurs iff the tweak is zero). */
            secp256k1_ge_neg(&pubtweak_neg, &pubtweak);
            secp256k1_gej_add_ge_var(&pubsharej, &pubsharej, &pubtweak_neg, NULL);
            secp256k1_ge_set_gej_var(&inv_data->pubshare, &pubsharej);
        }
        secp256k1_scalar_clear(&tweak);
        secp256k1_scalar_clear(&secshare_tweaked);
        return SECP256K1_CHILLDKG_UNKNOWN_FAULTY_PARTICIPANT_OR_COORDINATOR;
    }

    /* dkg_output: tweaked secshare, threshold public key (the tweaked
     * commitment to the secret), and the tweaked pubshares of all
     * participants. */
    secp256k1_scalar_get_b32(dkg_output->secshare32, &secshare_tweaked);
    secp256k1_chilldkg_point_save(dkg_output->thresh_pk33, &sum_coms_tweaked[0]);
    for (i = 0; i < n; i++) {
        if (i != participant_id) {
            /* We have computed our own pubshare already. */
            secp256k1_chilldkg_vss_pubshare(&pubsharej, sum_coms_tweaked, t, i);
            secp256k1_ge_set_gej_var(&pubshare_i, &pubsharej);
            secp256k1_chilldkg_point_save(dkg_output->pubshares33[i], &pubshare_i);
        } else {
            secp256k1_chilldkg_point_save(dkg_output->pubshares33[i], &pubshare);
        }
    }

    /* eq_input = u32be(t) || sum_coms.to_bytes() (UNTWEAKED commitment). */
    secp256k1_write_be32(eq_input, t);
    for (j = 0; j < t; j++) {
        secp256k1_chilldkg_point_save(eq_input + 4 + 33 * j, &sum_coms[j]);
    }

    secp256k1_scalar_clear(&tweak);
    secp256k1_scalar_clear(&secshare_tweaked);
    return SECP256K1_CHILLDKG_OK;
}

static void secp256k1_chilldkg_simplpedpop_assemble_sum_coms(secp256k1_ge *sum_coms, const secp256k1_ge *coms_to_secrets, const secp256k1_ge *sum_coms_to_nonconst_terms, size_t n, size_t t) {
    secp256k1_gej accj, termj;
    size_t i, j;

    /* Sum the commitments to the secrets. */
    secp256k1_gej_set_infinity(&accj);
    for (i = 0; i < n; i++) {
        secp256k1_gej_set_ge(&termj, &coms_to_secrets[i]);
        secp256k1_gej_add_var(&accj, &accj, &termj, NULL);
    }
    secp256k1_ge_set_gej_var(&sum_coms[0], &accj);
    for (j = 1; j < t; j++) {
        sum_coms[j] = sum_coms_to_nonconst_terms[j - 1];
    }
}

static secp256k1_chilldkg_fault secp256k1_chilldkg_simplpedpop_coordinator_step(const secp256k1_context *ctx, unsigned char *cmsg, secp256k1_chilldkg_simplpedpop_dkg_output *dkg_output, unsigned char *eq_input, uint32_t *fault_index, const unsigned char *const *pmsgs, uint32_t t, uint32_t n) {
    secp256k1_ge coms_to_secrets[SECP256K1_CHILLDKG_MAX_PARTICIPANTS];
    secp256k1_ge com[SECP256K1_CHILLDKG_MAX_PARTICIPANTS];
    secp256k1_ge sum_coms_to_nonconst_terms[SECP256K1_CHILLDKG_MAX_PARTICIPANTS - 1];
    secp256k1_ge sum_coms[SECP256K1_CHILLDKG_MAX_PARTICIPANTS];
    secp256k1_ge sum_coms_tweaked[SECP256K1_CHILLDKG_MAX_PARTICIPANTS];
    secp256k1_ge pubtweak;
    secp256k1_ge pubshare;
    secp256k1_gej pubsharej;
    secp256k1_scalar tweak;
    uint32_t i;
    size_t j;

    VERIFY_CHECK(ctx != NULL);
    ARG_CHECK(secp256k1_ecmult_gen_context_is_built(&ctx->ecmult_gen_ctx));
    ARG_CHECK(cmsg != NULL);
    ARG_CHECK(dkg_output != NULL);
    ARG_CHECK(eq_input != NULL);
    ARG_CHECK(pmsgs != NULL);

    if (t < 1 || n < 1 || t > n || n > SECP256K1_CHILLDKG_MAX_PARTICIPANTS) {
        return SECP256K1_CHILLDKG_INVALID_INPUT;
    }

    /* Parse the participant messages. A malformed message blames its sender.
     * The pops are not verified here; the participants verify them in step 2. */
    for (j = 1; j < t; j++) {
        secp256k1_ge_set_infinity(&sum_coms_to_nonconst_terms[j - 1]);
    }
    for (i = 0; i < n; i++) {
        for (j = 0; j < t; j++) {
            if (!secp256k1_chilldkg_point_load(&com[j], pmsgs[i] + 33 * j)) {
                *fault_index = i;
                return SECP256K1_CHILLDKG_FAULTY_PARTICIPANT;
            }
        }
        /* Echo the commitment to the secret ... */
        coms_to_secrets[i] = com[0];
        /* ... and sum the commitments to the non-constant terms (Pedersen,
         * Section 5.1 of "Non-Interactive and Information-Theoretic Secure
         * Verifiable Secret Sharing"). All inputs are public. */
        for (j = 1; j < t; j++) {
            secp256k1_gej aj, bj, sumj;
            secp256k1_gej_set_ge(&aj, &sum_coms_to_nonconst_terms[j - 1]);
            secp256k1_gej_set_ge(&bj, &com[j]);
            secp256k1_gej_add_var(&sumj, &aj, &bj, NULL);
            secp256k1_ge_set_gej_var(&sum_coms_to_nonconst_terms[j - 1], &sumj);
        }
        /* Collect the pops in the coordinator message. */
        memcpy(cmsg + 33 * n + 33 * (t - 1) + 64 * i, pmsgs[i] + 33 * t, 64);
    }

    /* cmsg = CoordinatorMsg(coms_to_secrets, sum_coms_to_nonconst_terms, pops).to_bytes() */
    for (i = 0; i < n; i++) {
        secp256k1_chilldkg_point_save(cmsg + 33 * i, &coms_to_secrets[i]);
    }
    for (j = 1; j < t; j++) {
        secp256k1_chilldkg_point_save(cmsg + 33 * n + 33 * (j - 1), &sum_coms_to_nonconst_terms[j - 1]);
    }

    secp256k1_chilldkg_simplpedpop_assemble_sum_coms(sum_coms, coms_to_secrets, sum_coms_to_nonconst_terms, n, t);

    if (!secp256k1_chilldkg_vss_invalid_taproot_commit(ctx, sum_coms_tweaked, &tweak, &pubtweak, sum_coms, t)) {
        /* A participant can make the sum of the commitments to the secrets
         * the point at infinity (the coordinator does not verify the pops),
         * or the tweak hash overflows (negligible probability). The faulty
         * participant cannot be identified without verifying the pops, which
         * is the participants' job in step 2. The reference raises a
         * non-protocol error here. */
        *fault_index = UINT32_MAX;
        secp256k1_scalar_clear(&tweak);
        return SECP256K1_CHILLDKG_UNKNOWN_FAULTY_PARTICIPANT_OR_COORDINATOR;
    }

    /* The coordinator's DKG output has no secshare (None in the reference). */
    memset(dkg_output->secshare32, 0, 32);
    secp256k1_chilldkg_point_save(dkg_output->thresh_pk33, &sum_coms_tweaked[0]);
    for (i = 0; i < n; i++) {
        secp256k1_chilldkg_vss_pubshare(&pubsharej, sum_coms_tweaked, t, i);
        secp256k1_ge_set_gej_var(&pubshare, &pubsharej);
        secp256k1_chilldkg_point_save(dkg_output->pubshares33[i], &pubshare);
    }

    /* eq_input = u32be(t) || sum_coms.to_bytes() (UNTWEAKED commitment). */
    secp256k1_write_be32(eq_input, t);
    for (j = 0; j < t; j++) {
        secp256k1_chilldkg_point_save(eq_input + 4 + 33 * j, &sum_coms[j]);
    }

    secp256k1_scalar_clear(&tweak);
    return SECP256K1_CHILLDKG_OK;
}

static secp256k1_chilldkg_fault secp256k1_chilldkg_simplpedpop_participant_investigate(const secp256k1_context *ctx, uint32_t *fault_index, const secp256k1_chilldkg_simplpedpop_participant_inv_data *inv_data, const secp256k1_ge *partial_pubshares, const secp256k1_scalar *partial_secshares) {
    const uint32_t n = inv_data->n;
    const uint32_t participant_id = inv_data->participant_id;
    secp256k1_gej sumj;
    secp256k1_ge sum;
    secp256k1_scalar secshare_sum;
    unsigned char buf33[33];
    unsigned char buf33b[33];
    uint32_t i;

    VERIFY_CHECK(ctx != NULL);
    ARG_CHECK(secp256k1_ecmult_gen_context_is_built(&ctx->ecmult_gen_ctx));

    /* If the sum of the partial pubshares is not equal to the pubshare, the
     * coordinator is faulty. */
    secp256k1_gej_set_infinity(&sumj);
    for (i = 0; i < n; i++) {
        secp256k1_gej termj;
        secp256k1_gej_set_ge(&termj, &partial_pubshares[i]);
        secp256k1_gej_add_var(&sumj, &sumj, &termj, NULL);
    }
    secp256k1_ge_set_gej_var(&sum, &sumj);
    secp256k1_chilldkg_point_save(buf33, &sum);
    secp256k1_chilldkg_point_save(buf33b, &inv_data->pubshare);
    if (secp256k1_memcmp_var(buf33, buf33b, 33) != 0) {
        return SECP256K1_CHILLDKG_FAULTY_COORDINATOR;
    }

    /* If the sum of the partial secshares is not equal to the secshare, the
     * coordinator is faulty (since the encryption is additively homomorphic,
     * this means that the sum of the *encrypted* partial secshares in the
     * investigation message was not equal to the encrypted secshare; the
     * reference signals this case with SecshareSumError, which the EncPedPop
     * layer translates to FaultyCoordinatorError). */
    secp256k1_scalar_set_int(&secshare_sum, 0);
    for (i = 0; i < n; i++) {
        secp256k1_scalar_add(&secshare_sum, &secshare_sum, &partial_secshares[i]);
    }
    /* We can declassify the result of the comparison because branching on it
     * only leaks whether the sum of the partial secshares equals the
     * secshare, which is exactly what the resulting public fault code
     * reveals. */
    {
        int sums_match = secp256k1_scalar_eq(&secshare_sum, &inv_data->secshare);
        secp256k1_declassify(ctx, &sums_match, sizeof(sums_match));
        if (!sums_match) {
            secp256k1_scalar_clear(&secshare_sum);
            return SECP256K1_CHILLDKG_FAULTY_COORDINATOR;
        }
    }
    secp256k1_scalar_clear(&secshare_sum);

    /* If some partial secshare does not match its partial pubshare, the
     * corresponding participant or the coordinator is faulty; for our own
     * index, only the coordinator can be at fault ("Coordinator fiddled with
     * the share from me to myself"). */
    for (i = 0; i < n; i++) {
        if (!secp256k1_chilldkg_vss_verify_secshare(ctx, &partial_secshares[i], &partial_pubshares[i])) {
            if (i != participant_id) {
                *fault_index = i;
                return SECP256K1_CHILLDKG_FAULTY_PARTICIPANT_OR_COORDINATOR;
            }
            return SECP256K1_CHILLDKG_FAULTY_COORDINATOR;
        }
    }

    /* We now know that the secshare matches the pubshare. This means that
     * this function was called even though no fault occurred (the reference
     * raises a RuntimeError in this case). */
    return SECP256K1_CHILLDKG_INVALID_INPUT;
}

#endif
