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

#ifndef SECP256K1_MODULE_FROST_KEYGEN_IMPL_H
#define SECP256K1_MODULE_FROST_KEYGEN_IMPL_H

#include <string.h>

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

#include "keygen.h"
#include "../../eckey.h"
#include "../../ecmult.h"
#include "../../util.h"

/* Initializes SHA256 with the BIP 445 "BIP0445/trusted/keygen" tag. */
static void secp256k1_frost_sha256_tagged_trusted_keygen(const secp256k1_hash_ctx *hash_ctx, secp256k1_sha256 *sha) {
    secp256k1_sha256_initialize_tagged(hash_ctx, sha, (const unsigned char *)"BIP0445/trusted/keygen", sizeof("BIP0445/trusted/keygen") - 1);
}

/* Derives the i-th (1 <= i <= t-1) trusted-dealer polynomial coefficient as
 *   a_i = tagged_hash("BIP0445/trusted/keygen", thresh_sk32 || ser32(i)).
 * Returns 0 if the hash output is zero or overflows the curve order (which can
 * only happen with negligible probability). */
static int secp256k1_frost_derive_coefficient(const secp256k1_hash_ctx *hash_ctx, secp256k1_scalar *out, const unsigned char *thresh_sk32, uint32_t i) {
    secp256k1_sha256 sha;
    unsigned char buf[4];
    unsigned char hash32[32];
    int overflow;

    secp256k1_frost_sha256_tagged_trusted_keygen(hash_ctx, &sha);
    secp256k1_sha256_write(hash_ctx, &sha, thresh_sk32, 32);
    secp256k1_write_be32(buf, i);
    secp256k1_sha256_write(hash_ctx, &sha, buf, sizeof(buf));
    secp256k1_sha256_finalize(hash_ctx, &sha, hash32);
    secp256k1_memclear_explicit(buf, sizeof(buf));
    secp256k1_scalar_set_b32(out, hash32, &overflow);
    secp256k1_memclear_explicit(hash32, sizeof(hash32));
    /* Bitwise (not short-circuiting) so that this does not branch on the
     * secret-derived overflow; the caller declassifies the return value. */
    return (!overflow) & (!secp256k1_scalar_is_zero(out));
}

static const unsigned char secp256k1_frost_tweak_cache_magic[4] = { 0x8d, 0x86, 0xb5, 0x01 };

/* A tweak cache consists of
 * - 4 byte magic set during initialization to allow detecting an uninitialized
 *   object.
 * - 64 byte untweaked threshold public key (Q_0)
 * - 64 byte current (tweaked) threshold public key (Q)
 * - 32 byte tweak accumulator tacc
 * - 1 byte gacc_parity (1 if the tweak factor accumulator gacc is -1)
 */
/* Requires that cache_i->thresh_pk and cache_i->pk are not infinity. */
static void secp256k1_frost_tweak_cache_save(secp256k1_frost_tweak_cache *cache, const secp256k1_frost_tweak_cache_internal *cache_i) {
    unsigned char *ptr = cache->data;

    VERIFY_CHECK(!secp256k1_ge_is_infinity(&cache_i->thresh_pk));
    VERIFY_CHECK(!secp256k1_ge_is_infinity(&cache_i->pk));
    memcpy(ptr, secp256k1_frost_tweak_cache_magic, 4);
    ptr += 4;
    secp256k1_ge_to_bytes(ptr, &cache_i->thresh_pk);
    ptr += 64;
    secp256k1_ge_to_bytes(ptr, &cache_i->pk);
    ptr += 64;
    secp256k1_scalar_get_b32(ptr, &cache_i->tacc);
    ptr += 32;
    *ptr = cache_i->gacc_parity & 1;
}

static int secp256k1_frost_tweak_cache_load(const secp256k1_context* ctx, secp256k1_frost_tweak_cache_internal *cache_i, const secp256k1_frost_tweak_cache *cache) {
    const unsigned char *ptr = cache->data;

    ARG_CHECK(secp256k1_memcmp_var(ptr, secp256k1_frost_tweak_cache_magic, 4) == 0);
    ptr += 4;
    secp256k1_ge_from_bytes(&cache_i->thresh_pk, ptr);
    ptr += 64;
    secp256k1_ge_from_bytes(&cache_i->pk, ptr);
    ptr += 64;
    secp256k1_scalar_set_b32(&cache_i->tacc, ptr, NULL);
    ptr += 32;
    cache_i->gacc_parity = *ptr & 1;
    return 1;
}

/* Returns 1 if the ids array contains no duplicates, 0 otherwise. */
static int secp256k1_frost_ids_are_valid(const uint32_t *ids, size_t n_ids) {
    size_t i, j;

    for (i = 0; i < n_ids; i++) {
        for (j = i + 1; j < n_ids; j++) {
            if (ids[i] == ids[j]) {
                return 0;
            }
        }
    }
    return 1;
}

static int secp256k1_frost_derive_interpolating_value(secp256k1_scalar *out, const uint32_t *ids, size_t n_ids, uint32_t my_id) {
    secp256k1_scalar num, deno, one;
    size_t i;
    int found = 0;

    if (!secp256k1_frost_ids_are_valid(ids, n_ids)) {
        return 0;
    }
    for (i = 0; i < n_ids; i++) {
        if (ids[i] == my_id) {
            found = 1;
            break;
        }
    }
    if (!found) {
        return 0;
    }

    /* lambda = product_{j != i} x_j / (x_j - x_i), where the x-coordinate of
     * an identifier id is id + 1. The +1 is added in scalar arithmetic because
     * it would overflow for id == UINT32_MAX, which the BIP 445
     * derive_interpolating_value algorithm permits (0 <= id < 2^32). */
    secp256k1_scalar_set_int(&num, 1);
    secp256k1_scalar_set_int(&deno, 1);
    secp256k1_scalar_set_int(&one, 1);
    for (i = 0; i < n_ids; i++) {
        secp256k1_scalar id_j, x_j, term;
        if (ids[i] == my_id) {
            continue;
        }
        secp256k1_scalar_set_int(&id_j, ids[i]);
        /* x_j = id_j + 1 */
        secp256k1_scalar_add(&x_j, &id_j, &one);
        secp256k1_scalar_mul(&num, &num, &x_j);
        /* x_j - x_i = (id_j + 1) - (my_id + 1) = id_j - my_id */
        secp256k1_scalar_set_int(&term, my_id);
        secp256k1_scalar_negate(&term, &term);
        secp256k1_scalar_add(&term, &term, &id_j);
        secp256k1_scalar_mul(&deno, &deno, &term);
    }
    /* deno != 0 because the ids are distinct */
    VERIFY_CHECK(!secp256k1_scalar_is_zero(&deno));
    secp256k1_scalar_inverse_var(&deno, &deno);
    secp256k1_scalar_mul(out, &num, &deno);
    return 1;
}

static int secp256k1_frost_derive_pubshare_at(secp256k1_gej *out, const uint32_t *ids, const secp256k1_ge *pubshares, size_t n_ids, const secp256k1_scalar *x) {
    size_t i, j;

    if (!secp256k1_frost_ids_are_valid(ids, n_ids)) {
        return 0;
    }

    /* out = sum_i lambda_i * pubshares[i] with
     * lambda_i = product_{j != i} (x - id_j) / (id_i - id_j). */
    secp256k1_gej_set_infinity(out);
    for (i = 0; i < n_ids; i++) {
        secp256k1_scalar num, deno, lambda, id_i;
        secp256k1_gej term;

        secp256k1_scalar_set_int(&id_i, ids[i]);
        secp256k1_scalar_set_int(&num, 1);
        secp256k1_scalar_set_int(&deno, 1);
        for (j = 0; j < n_ids; j++) {
            secp256k1_scalar id_j;
            if (j == i) {
                continue;
            }
            secp256k1_scalar_set_int(&id_j, ids[j]);
            /* num *= x - id_j */
            secp256k1_scalar_negate(&id_j, &id_j);
            secp256k1_scalar_add(&id_j, &id_j, x);
            secp256k1_scalar_mul(&num, &num, &id_j);
            /* deno *= id_i - id_j */
            secp256k1_scalar_set_int(&id_j, ids[j]);
            secp256k1_scalar_negate(&id_j, &id_j);
            secp256k1_scalar_add(&id_j, &id_j, &id_i);
            secp256k1_scalar_mul(&deno, &deno, &id_j);
        }
        /* deno != 0 because the ids are distinct */
        VERIFY_CHECK(!secp256k1_scalar_is_zero(&deno));
        secp256k1_scalar_inverse_var(&deno, &deno);
        secp256k1_scalar_mul(&lambda, &num, &deno);

        VERIFY_CHECK(!secp256k1_ge_is_infinity(&pubshares[i]));
        secp256k1_gej_set_ge(&term, &pubshares[i]);
        secp256k1_ecmult(&term, &term, &lambda, NULL);
        secp256k1_gej_add_var(out, out, &term, NULL);
    }
    /* The result may be the point at infinity; the caller must handle it. */
    return 1;
}

static int secp256k1_frost_derive_thresh_pubkey(secp256k1_gej *out, const uint32_t *ids, const secp256k1_ge *pubshares, size_t n_ids) {
    /* Evaluating in identifier space at identifier -1 corresponds to
     * evaluating the polynomial at x-coordinate 0. */
    secp256k1_scalar minus_one;

    secp256k1_scalar_set_int(&minus_one, 1);
    secp256k1_scalar_negate(&minus_one, &minus_one);
    if (!secp256k1_frost_derive_pubshare_at(out, ids, pubshares, n_ids, &minus_one)) {
        return 0;
    }
    /* The threshold public key must not be the point at infinity. */
    return !secp256k1_gej_is_infinity(out);
}

int secp256k1_frost_trusted_dealer_keygen(const secp256k1_context *ctx, unsigned char *secshares32, secp256k1_pubkey *thresh_pk, secp256k1_pubkey *pubshares, size_t n_participants, uint32_t threshold, const unsigned char *threshold_seckey32) {
    /* At most SECP256K1_FROST_MAX_PARTICIPANTS - 1 coefficients */
    secp256k1_scalar coeffs[SECP256K1_FROST_MAX_PARTICIPANTS - 1];
    secp256k1_scalar secret, share, x;
    secp256k1_gej pkj;
    secp256k1_ge pk;
    size_t i, j;
    int overflow;
    int ret = 0;

    VERIFY_CHECK(ctx != NULL);
    ARG_CHECK(secshares32 != NULL);
    ARG_CHECK(thresh_pk != NULL);
    ARG_CHECK(pubshares != NULL);
    ARG_CHECK(threshold_seckey32 != NULL);
    ARG_CHECK(secp256k1_ecmult_gen_context_is_built(&ctx->ecmult_gen_ctx));

    if (n_participants < 1 || n_participants > SECP256K1_FROST_MAX_PARTICIPANTS) {
        return 0;
    }
    if (threshold < 1 || (size_t)threshold > n_participants) {
        return 0;
    }
    /* Ensure that the outputs are unusable if this function fails. */
    memset(secshares32, 0, n_participants * 32);
    memset(thresh_pk, 0, sizeof(*thresh_pk));
    memset(pubshares, 0, n_participants * sizeof(*pubshares));

    secp256k1_scalar_set_b32(&secret, threshold_seckey32, &overflow);
    overflow |= secp256k1_scalar_is_zero(&secret);
    /* We can declassify because branching on overflow only leaks whether the
     * threshold secret key is valid, which is not secret. */
    secp256k1_declassify(ctx, &overflow, sizeof(overflow));
    if (overflow) {
        goto cleanup;
    }

    /* The coefficients a_1, ..., a_{t-1} (a_1 is the highest-degree term) are
     * derived deterministically from the threshold secret key. */
    for (i = 0; i + 1 < threshold; i++) {
        /* A failure can only occur with negligible probability, so its
         * occurrence is not secret. */
        int coeff_ok = secp256k1_frost_derive_coefficient(secp256k1_get_hash_context(ctx), &coeffs[i], threshold_seckey32, (uint32_t)(i + 1));
        secp256k1_declassify(ctx, &coeff_ok, sizeof(coeff_ok));
        if (!coeff_ok) {
            goto cleanup;
        }
    }

    /* The secret share of participant id i is p(i + 1), where the polynomial
     * p(x) = a_1*x^(t-1) + ... + a_{t-1}*x + secret is evaluated with Horner's
     * method. */
    for (i = 0; i < n_participants; i++) {
        secp256k1_scalar_set_int(&x, (unsigned int)(i + 1));
        secp256k1_scalar_set_int(&share, 0);
        for (j = 0; j + 1 < threshold; j++) {
            secp256k1_scalar_mul(&share, &share, &x);
            secp256k1_scalar_add(&share, &share, &coeffs[j]);
        }
        secp256k1_scalar_mul(&share, &share, &x);
        secp256k1_scalar_add(&share, &share, &secret);
        /* A zero share can only occur with negligible probability, so its
         * occurrence is not secret. */
        {
            int share_is_zero = secp256k1_scalar_is_zero(&share);
            secp256k1_declassify(ctx, &share_is_zero, sizeof(share_is_zero));
            if (share_is_zero) {
                goto cleanup;
            }
        }
        secp256k1_scalar_get_b32(&secshares32[32 * i], &share);

        secp256k1_ecmult_gen_gej(&ctx->ecmult_gen_ctx, &pkj, &share);
        secp256k1_ge_set_gej(&pk, &pkj);
        secp256k1_pubkey_save(&pubshares[i], &pk);
    }

    /* The threshold public key is the plain (parity-preserving) encoding of
     * threshold_seckey32*G. */
    secp256k1_ecmult_gen_gej(&ctx->ecmult_gen_ctx, &pkj, &secret);
    secp256k1_ge_set_gej(&pk, &pkj);
    secp256k1_pubkey_save(thresh_pk, &pk);
    ret = 1;

cleanup:
    secp256k1_scalar_clear(&secret);
    secp256k1_scalar_clear(&share);
    secp256k1_scalar_clear(&x);
    for (i = 0; i + 1 < threshold; i++) {
        secp256k1_scalar_clear(&coeffs[i]);
    }
    secp256k1_ge_clear(&pk);
    secp256k1_gej_clear(&pkj);
    return ret;
}

int secp256k1_frost_threshold_info_validate(const secp256k1_context *ctx, const secp256k1_pubkey *thresh_pk, const secp256k1_pubkey *pubshares, size_t n_participants, uint32_t threshold) {
    secp256k1_ge base_points[SECP256K1_FROST_MAX_PARTICIPANTS];
    uint32_t base_ids[SECP256K1_FROST_MAX_PARTICIPANTS];
    secp256k1_ge pk, expected, derived;
    secp256k1_gej result;
    size_t i;

    VERIFY_CHECK(ctx != NULL);
    ARG_CHECK(thresh_pk != NULL);
    ARG_CHECK(pubshares != NULL);

    if (n_participants < 1 || n_participants > SECP256K1_FROST_MAX_PARTICIPANTS) {
        return 0;
    }
    if (threshold < 1 || (size_t)threshold > n_participants) {
        return 0;
    }

    /* A corrupted or uninitialized pubkey fails in pubkey_load. */
    if (!secp256k1_pubkey_load(ctx, &pk, thresh_pk)) {
        return 0;
    }

    /* The base set for polynomial interpolation consists of the first t
     * pubshares (identifiers 0, ..., t-1). */
    for (i = 0; i < threshold; i++) {
        base_ids[i] = (uint32_t)i;
        if (!secp256k1_pubkey_load(ctx, &base_points[i], &pubshares[i])) {
            return 0;
        }
    }

    /* Every remaining pubshare must lie on the same polynomial as the base
     * set. */
    for (i = threshold; i < n_participants; i++) {
        secp256k1_scalar x;
        if (!secp256k1_pubkey_load(ctx, &expected, &pubshares[i])) {
            return 0;
        }
        secp256k1_scalar_set_int(&x, (unsigned int)i);
        if (!secp256k1_frost_derive_pubshare_at(&result, base_ids, base_points, threshold, &x)) {
            return 0;
        }
        if (secp256k1_gej_is_infinity(&result)) {
            return 0;
        }
        secp256k1_ge_set_gej_var(&derived, &result);
        if (!secp256k1_ge_eq_var(&expected, &derived)) {
            return 0;
        }
    }

    /* The public key derived from the base set must match the provided
     * threshold public key. */
    if (!secp256k1_frost_derive_thresh_pubkey(&result, base_ids, base_points, threshold)) {
        return 0;
    }
    secp256k1_ge_set_gej_var(&derived, &result);
    return secp256k1_ge_eq_var(&pk, &derived);
}

int secp256k1_frost_tweak_cache_init(const secp256k1_context *ctx, secp256k1_frost_tweak_cache *cache, const secp256k1_pubkey *thresh_pk) {
    secp256k1_frost_tweak_cache_internal cache_i = { 0 };

    VERIFY_CHECK(ctx != NULL);
    ARG_CHECK(cache != NULL);
    memset(cache, 0, sizeof(*cache));
    ARG_CHECK(thresh_pk != NULL);

    if (!secp256k1_pubkey_load(ctx, &cache_i.pk, thresh_pk)) {
        return 0;
    }
    /* Q = thresh_pk (parity is preserved), gacc = 1, tacc = 0. */
    cache_i.thresh_pk = cache_i.pk;
    secp256k1_scalar_set_int(&cache_i.tacc, 0);
    cache_i.gacc_parity = 0;
    secp256k1_frost_tweak_cache_save(cache, &cache_i);
    return 1;
}

int secp256k1_frost_tweaked_pubkey_get(const secp256k1_context *ctx, secp256k1_xonly_pubkey *tweaked_pk, const secp256k1_frost_tweak_cache *cache) {
    secp256k1_frost_tweak_cache_internal cache_i;

    VERIFY_CHECK(ctx != NULL);
    ARG_CHECK(tweaked_pk != NULL);
    memset(tweaked_pk, 0, sizeof(*tweaked_pk));
    ARG_CHECK(cache != NULL);

    if (!secp256k1_frost_tweak_cache_load(ctx, &cache_i, cache)) {
        return 0;
    }
    /* x-only public keys carry the even-Y representative (required for BIP340
     * verification), so the point is normalized to even Y before saving. */
    secp256k1_extrakeys_ge_even_y(&cache_i.pk);
    secp256k1_xonly_pubkey_save(tweaked_pk, &cache_i.pk);
    return 1;
}

/* Implements ApplyTweak of BIP 445: computes Q' = g*Q + tweak*G where g = -1
 * iff xonly is set and Q has odd Y, updates tacc = tweak + g*tacc and
 * gacc_parity accordingly, and saves the result to the cache. On failure
 * (tweak out of range or Q' is the point at infinity) the cache is left
 * unchanged. */
static int secp256k1_frost_pubkey_tweak_add_internal(const secp256k1_context *ctx, secp256k1_xonly_pubkey *tweaked_pk, secp256k1_frost_tweak_cache *cache, const unsigned char *tweak32, int xonly) {
    secp256k1_frost_tweak_cache_internal cache_i;
    secp256k1_scalar tweak;
    int overflow = 0;

    VERIFY_CHECK(ctx != NULL);
    if (tweaked_pk != NULL) {
        memset(tweaked_pk, 0, sizeof(*tweaked_pk));
    }
    ARG_CHECK(cache != NULL);
    ARG_CHECK(tweak32 != NULL);

    if (!secp256k1_frost_tweak_cache_load(ctx, &cache_i, cache)) {
        return 0;
    }
    secp256k1_scalar_set_b32(&tweak, tweak32, &overflow);
    if (overflow) {
        return 0;
    }
    /* For xonly tweaking, g = -1 iff the current public key has odd Y. All
     * modifications happen in the local cache_i and are only committed on
     * success. */
    if (xonly && secp256k1_extrakeys_ge_even_y(&cache_i.pk)) {
        cache_i.gacc_parity ^= 1;
        secp256k1_scalar_negate(&cache_i.tacc, &cache_i.tacc);
    }
    secp256k1_scalar_add(&cache_i.tacc, &cache_i.tacc, &tweak);
    if (!secp256k1_eckey_pubkey_tweak_add(&cache_i.pk, &tweak)) {
        return 0;
    }
    /* eckey_pubkey_tweak_add fails if the result is the point at infinity */
    VERIFY_CHECK(!secp256k1_ge_is_infinity(&cache_i.pk));
    secp256k1_frost_tweak_cache_save(cache, &cache_i);
    if (tweaked_pk != NULL) {
        secp256k1_ge pk_even = cache_i.pk;
        /* x-only public keys carry the even-Y representative (required for
         * BIP340 verification) */
        secp256k1_fe_normalize_var(&pk_even.y);
        secp256k1_extrakeys_ge_even_y(&pk_even);
        secp256k1_xonly_pubkey_save(tweaked_pk, &pk_even);
    }
    return 1;
}

int secp256k1_frost_pubkey_ec_tweak_add(const secp256k1_context *ctx, secp256k1_xonly_pubkey *tweaked_pk, secp256k1_frost_tweak_cache *cache, const unsigned char *tweak32) {
    return secp256k1_frost_pubkey_tweak_add_internal(ctx, tweaked_pk, cache, tweak32, 0);
}

int secp256k1_frost_pubkey_xonly_tweak_add(const secp256k1_context *ctx, secp256k1_xonly_pubkey *tweaked_pk, secp256k1_frost_tweak_cache *cache, const unsigned char *tweak32) {
    return secp256k1_frost_pubkey_tweak_add_internal(ctx, tweaked_pk, cache, tweak32, 1);
}

#endif
