#ifndef TOPRF_UPDATE_H
#define TOPRF_UPDATE_H

/**
 * @file toprf-update.h
 * @brief TOPRF Update
 *
 * SPDX-FileCopyrightText: 2025, Marsiske Stefan
 * SPDX-License-Identifier: LGPL-3.0-or-later
 *
 * This file defines the structures, types, and functions for updating a threshold OPRF
 * instantiation.
 */

#include <stdint.h>
#include <stdio.h>
#include <sodium.h>
#include "dkg.h"
#include "toprf.h"

#define noise_xk_handshake1_SIZE 48UL
#define noise_xk_handshake2_SIZE 48UL
#define noise_xk_handshake3_SIZE 64UL
#define toprf_update_noise_key_SIZE (32UL)
#define toprf_update_sessionid_SIZE 32U
#define toprf_update_commitment_HASHBYTES 32U
#define toprf_update_encrypted_shares_SIZE (TOPRF_Share_BYTES * 2)
#define toprf_keyid_SIZE 32U

typedef DKG_Message TOPRF_Update_Message;

// TODO: refactor this, as it's the same in tp-dkg and stp-dkg
/**
 * @struct TOPRF_Update_Cheater
 * @brief Records information about protocol violations
 *
 * This struct is used to communicate and track detailed information about
 * detected violations of the protocol rules by participants.
 *
 * @var TOPRF_Update_Cheater::step The protocol step in which the violation occurred
 * @var TOPRF_Update_Cheater::error Error code specifying the type of violation
 * @var TOPRF_Update_Cheater::peer Index of the peer that caused the violation
 * @var TOPRF_Update_Cheater::other_peer Optional index of the peer that reported the
 *  violation (set to 0xfe if unused)
 * @var TOPRF_Update_Cheater::invalid_index
 */
typedef struct {
  int step;
  int error;
  uint8_t peer;
  uint8_t other_peer;
  int invalid_index;
} TOPRF_Update_Cheater;

/**
 * @enum TOPRF_Update_Err
 * @brief Error codes for the TOPRF-Update protocol
 *
 * These error codes provide detailed information about various failure
 * conditions that can occur during the TOPRF update protocol execution.
 */
typedef enum {
  TOPRF_Update_Err_OK = 0,
  TOPRF_Update_Err_ISize,
  TOPRF_Update_Err_OSize,
  TOPRF_Update_Err_OOB,
  TOPRF_Update_Err_Send,
  TOPRF_Update_Err_CheatersFound,
  TOPRF_Update_Err_CheatersFull,
  TOPRF_Update_Err_InvSessionID,
  TOPRF_Update_Err_VSSShare,
  TOPRF_Update_Err_VSSCommit,
  TOPRF_Update_Err_Noise,
  TOPRF_Update_Err_NoiseEncrypt,
  TOPRF_Update_Err_NoiseDecrypt,
  TOPRF_Update_Err_HMac,
  TOPRF_Update_Err_NoSubVSPSFail,
  TOPRF_Update_Err_NotEnoughDealers,
  TOPRF_Update_Err_TooManyCheaters,
  TOPRF_Update_Err_DKGFinish,
  TOPRF_Update_Err_FTMULTStep1,
  TOPRF_Update_Err_FTMULTZKCommitments,
  TOPRF_Update_Err_InvPoint,
  TOPRF_Update_Err_CommmitmentsMismatch,
  TOPRF_Update_Err_Proto,
  TOPRF_Update_Err_BadReconstruct,
  TOPRF_Update_Err_Reconstruct,
  TOPRF_Update_Err_BroadcastEnv = 32,
  TOPRF_Update_Err_Env = 64
} TOPRF_Update_Err;

/**
 * @enum TOPRF_Update_STP_Steps
 * @brief TOPRF-Update protocol steps for the STP
 *
 * These values define the sequential steps that the STP
 * follows during the execution of the TOPRF update protocol.
 */
typedef enum {
  TOPRF_Update_STP_Broadcast_NPKs,
  TOPRF_Update_STP_Route_Noise_Handshakes1,
  TOPRF_Update_STP_Route_Noise_Handshakes2,
  TOPRF_Update_STP_Broadcast_DKG_Hash_Commitments,
  TOPRF_Update_STP_Broadcast_DKG_Commitments,
  TOPRF_Update_STP_Route_Encrypted_Shares,
  TOPRF_Update_STP_Broadcast_Complaints,
  TOPRF_Update_STP_Broadcast_DKG_Defenses,
  TOPRF_Update_STP_Broadcast_DKG_Transcripts,
  TOPRF_Update_STP_Route_Mult_Step1,
  TOPRF_Update_STP_Broadcast_Mult_Commitments,
  TOPRF_Update_STP_Route_Encrypted_Mult_Shares,
  TOPRF_Update_STP_Broadcast_Mult_Complaints,
  TOPRF_Update_STP_Broadcast_Mult_Defenses,
  TOPRF_Update_STP_Broadcast_Reconst_Mult_Shares,
  TOPRF_Update_STP_Route_ZK_Challenge_Commitments,
  TOPRF_Update_STP_Route_ZK_commitments,
  TOPRF_Update_STP_Broadcast_ZK_nonces,
  TOPRF_Update_STP_Broadcast_ZK_Proofs,
  TOPRF_Update_STP_Broadcast_ZK_Disclosures,
  TOPRF_Update_STP_Broadcast_Mult_Ci,
  TOPRF_Update_STP_Broadcast_VSPS_Disclosures,
  TOPRF_Update_STP_Reconstruct_Delta,
  TOPRF_Update_STP_Done
} TOPRF_Update_STP_Steps;

/**
 * @struct TOPRF_Update_STPState
 * @brief State of the STP during protocol execution
 *
 * This struct maintains the state of the STP during the execution of
 * the TOPRF Update protocol.
 * Some fields in this struct are internal variables and should not
 * be used. The following fields are useful and can be accessed by
 * users of the API:
 *
 * @var TOPRF_Update_STPState::n The total number of peers participating
 *      in this protocol
 *
 * @var TOPRF_Update_STPState::t The threshold necessary to use shared
 *      secret generated by this DKG
 *
 * @var TOPRF_Update_STPState::step the current step in the protocol
 *
 * @var TOPRF_Update_STPState::prev the previous step in the protocol
 *
 * @var TOPRF_Update_STPState::cheaters List of detected cheaters and
 *      protocol violators at the end of a failed protocol run
 *
 * @var TOPRF_Update_STPState::delta The final public delta value, which
 *      can be applied to previous OPRF results with the old key to update
 *      those results to the new key
 *
 * @var TOPRF_Update_STPState::k0p_final_commitments The final
 *      commitments to the shares of the updated key.
 */

// TODO: consider making this struct private
typedef struct {
  TOPRF_Update_STP_Steps step;
  TOPRF_Update_STP_Steps prev;
  uint8_t sessionid[dkg_sessionid_SIZE];
  uint8_t n;
  uint8_t t;
  uint8_t sig_pk[crypto_sign_PUBLICKEYBYTES];
  uint8_t sig_sk[crypto_sign_SECRETKEYBYTES];
  uint64_t *last_ts;
  uint64_t ts_epsilon;
  const uint8_t (*sig_pks)[][crypto_sign_PUBLICKEYBYTES];

  uint8_t (*p_commitments_hashes)[][toprf_update_commitment_HASHBYTES];
  uint8_t (*p_share_macs)[][crypto_auth_hmacsha256_BYTES];
  uint8_t (*p_commitments)[][crypto_core_ristretto255_BYTES];
  uint16_t p_complaints_len;
  uint16_t *p_complaints;
  uint16_t y2_complaints_len;
  uint16_t *y2_complaints;

  uint8_t (*kc0_commitments)[][crypto_core_ristretto255_BYTES];
  uint8_t (*k0p_commitments)[][crypto_core_ristretto255_BYTES];
  uint8_t (*zk_challenge_commitments)[][3][crypto_scalarmult_ristretto255_SCALARBYTES];
  uint8_t (*zk_challenge_e_i)[][crypto_scalarmult_ristretto255_SCALARBYTES];

  size_t cheater_len;
  TOPRF_Update_Cheater (*cheaters)[];
  size_t cheater_max;
  uint8_t (*k0p_final_commitments)[][crypto_scalarmult_ristretto255_BYTES];
  uint8_t delta[crypto_scalarmult_ristretto255_BYTES];
  crypto_generichash_state transcript_state;
  uint8_t transcript[crypto_generichash_BYTES];
} TOPRF_Update_STPState;

/**
 * @brief Gets the size of a TOPRF_Update_STPState struct
 *
 * Returns the memory size required for the STPState structure.
 * This is useful when dynamically allocating memory for this
 * structure. If you do, note that the struct must be aligned at
 * a 32-byte boundary for proper operation.
 * The python wrapper, pyoprf, handles this like this:
 * @code
 * import ctypes
 * ctypes.create_string_buffer(liboprf.toprf_update_stpstate_size()+32)
 * @endcode
 *
 * @return Size of the TOPRF_Update_STPState structure in bytes
 */
size_t toprf_update_stpstate_size(void);

/**
 * @brief Gets the value of `n` (number of participants) from the STP state
 *
 * @param[in] ctx Pointer to an initialized TOPRF_Update_STPState struct
 *
 * @return The number of participants in the protocol
 */
uint8_t toprf_update_stpstate_n(const TOPRF_Update_STPState *ctx);

/**
 * @brief Gets the value of `t` (threshold) from the STP state
 *
 * @param[in] ctx Pointer to an initialized TOPRF_Update_STPState struct
 *
 * @return The threshold value used in the protocol
 */
uint8_t toprf_update_stpstate_t(const TOPRF_Update_STPState *ctx);

/**
 * @brief Gets the number of cheaters discovered so far from the STP state
 *
 * @param[in] ctx Pointer to an initialized TOPRF_Update_STPState struct
 *
 * @return The count of detected cheaters and protocol violators
 */
size_t toprf_update_stpstate_cheater_len(const TOPRF_Update_STPState *ctx);

/**
 * @brief Gets the session ID from the STP state
 *
 * @param[in] ctx Pointer to an initialized TOPRF_Update_STPState struct
 *
 * @return Pointer to the session ID byte array
 */
const uint8_t *toprf_update_stpstate_sessionid(const TOPRF_Update_STPState *ctx);

/**
 * @brief Gets the final result, the public delta value from the STP state
 *
 * This function returns the public delta value which can be applied to
 * previous OPRF results to update them to use the new key. It should only
 * be used after the protocol has finished successfully.
 *
 * @param[in] ctx Pointer to an initialized TOPRF_Update_STPState struct
 *
 * @return Pointer to the delta value byte array, or NULL if protocol
 *         not complete
 */
const uint8_t *toprf_update_stpstate_delta(const TOPRF_Update_STPState *ctx);

/**
 * @brief Gets the commitments for the shares of the new key from the STP
 *        state
 *
 * This function returns the public commitments to the shares of the new key.
 * It should only be used after the protocol has finished successfully.
 *
 * @param[in] ctx Pointer to an initialized TOPRF_Update_STPState struct
 *
 * @return Pointer to the commitments byte array, or NULL if protocol
 *         not complete
 */
const uint8_t *toprf_update_stpstate_commitments(const TOPRF_Update_STPState *ctx);

/**
 * @brief Gets the current step of the STP in the protocol from the STP state
 *
 * @param[in] ctx Pointer to an initialized TOPRF_Update_STPState structure
 *
 * @return The current protocol step as an integer
 */
int toprf_update_stpstate_step(const TOPRF_Update_STPState *ctx);

#define toprfupdate_stp_start_msg_SIZE (sizeof(TOPRF_Update_Message) + crypto_generichash_BYTES /*dst*/ \
                                        + toprf_keyid_SIZE + crypto_sign_PUBLICKEYBYTES)

/**
 * @brief Starts a new execution of the STP TOPRF update protocol
 *
 * This function initializes the state of the STP and creates an initial
 * message containing the parameters for the peers to join the protocol.
 *
 * @param[in] ctx Pointer to a STP_DKG_STPState struct. This struct
 *            will be initialized by this function
 * @param[in] ts_epsilon Maximum allowed message age in seconds before
 *            it is considered stale and rejected. This value is used to
 *            prevent replay attacks and enforce freshness. For small,
 *            local setups (e.g., 2-out-of-3 participants), values as low
 *            as 2–3 seconds may suffice. For large-scale deployments
 *            (e.g., 126-out-of-127), this may need to be increased to
 *            several hours
 * @param[in] n Number of peers participating in this execution
 * @param[in] t Threshold necessary to use the results of this update.
 *            This value should be the same as it was for the previous
 *            key being shared
 * @param[in] proto_name An array of bytes used as a domain seperation tag
 *            (DST). Set it to the name of your application
 * @param[in] proto_name_len The size of the array `proto_name`, to allow
 *            non-zero terminated DSTs
 * @param[in] keyid The ID of the key to be updated. Must be exactly
 *            toprf_keyid_SIZE bytes long
 * @param[in] sig_pks Pointer to a (n+1)-element array of signing public
 *            keys. The STP's public key must be at index 0. The rest of the
 *            items must be in order
 * @param[in] ltssk Long-term signing private key of the STP
 * @param[in] msg0_len Size of memory allocated for the `msg0` parameter. It
 *            should be tpdkg_msg0_SIZE
 * @param[out] msg0 Message to be sent to all peers to initialize them
 *
 * @return 0 on success, non-zero on error
 */
int toprf_update_start_stp(TOPRF_Update_STPState *ctx, const uint64_t ts_epsilon,
                           const uint8_t n, const uint8_t t,
                           const char *proto_name, const size_t proto_name_len,
                           const uint8_t keyid[toprf_keyid_SIZE],
                           const uint8_t (*sig_pks)[][crypto_sign_PUBLICKEYBYTES],
                           const uint8_t ltssk[crypto_sign_SECRETKEYBYTES],
                           const size_t msg0_len, TOPRF_Update_Message *msg0);

/**
 * @brief Sets variable-sized buffers in the TOPRF_Update_STPState structure
 *
 * A number of buffers are needed in the STP state that depend on the `n`
 * (number of participants) and `t` (threshold) parameters.
 *
 * These can be allocated on the stack as follows:
 * @code
 *  uint16_t stp_p_complaints[n*n];
 *  uint16_t stp_y2_complaints[n*n];
 *  uint64_t last_ts[n];
 *
 *  uint8_t stp_p_commitments_hashes[n][toprf_update_commitment_HASHBYTES];
 *  uint8_t stp_p_share_macs[n*n][crypto_auth_hmacsha256_BYTES];
 *  uint8_t stp_p_commitments[n*n][crypto_core_ristretto255_BYTES];
 *
 *  uint8_t stp_k0p_commitments[dealers*(n+1)][crypto_core_ristretto255_BYTES];
 *  uint8_t stp_zk_challenge_commitments[dealers][3][crypto_scalarmult_ristretto255_SCALARBYTES];
 *  uint8_t stp_zk_challenge_e_i[dealers][crypto_scalarmult_ristretto255_SCALARBYTES];
 *
 *  uint8_t k0p_final_commitments[n][crypto_scalarmult_ristretto255_BYTES];
 *  TOPRF_Update_Cheater stp_cheaters[t*t - 1];
 *  toprf_update_stp_set_bufs(&stp,
 *                            stp_p_complaints,
 *                            stp_y2_complaints,
 *                            &stp_cheaters,
 *                            sizeof(stp_cheaters) / sizeof(TOPRF_Update_Cheater),
 *                            &stp_p_commitments_hashes,
 *                            &stp_p_commitments,
 *                            &k0_commitments,
 *                            &stp_k0p_commitments,
 *                            &stp_zk_challenge_commitments,
 *                            &stp_zk_challenge_e_i,
 *                            &k0p_final_commitments,
 *                            last_ts);
 * @endcode
 *
 * @param[in] ctx Pointer to initialized TOPRF_Update_STPState structure
 * @param[in] p_complaints Buffer for polynomial complaints
 * @param[in] cheaters List of detected cheaters and protocol violators at
 *            the end of a failed protocol run
 * @param[in] cheater_max Maximum number of cheat attempts to be recorded.
 *            Normally, the maximum number of cheaters is `t * t - 1`, where
 *            `t` is the threshold parameter. It should be provided as
 *            (sizeof(cheaters) / sizeof(TP_DKG_Cheater))
 * @param[in] p_commitments_hashes Buffer for polynomial commitment hashes
 * @param[in] p_share_macs Buffer for share Message Authentication Codes
 *            (MACs)
 * @param[in] p_commitments Buffer for polynomial commitments
 * @param[in] last_ts Pointer to a list of last timestamps
 */
void toprf_update_stp_set_bufs(TOPRF_Update_STPState *ctx,
                               uint16_t p_complaints[],
                               uint16_t y2_complaint[],
                               TOPRF_Update_Cheater (*cheaters)[], const size_t cheater_max,
                               uint8_t (*p_commitments_hashes)[][toprf_update_commitment_HASHBYTES],
                               uint8_t (*p_share_macs)[][crypto_auth_hmacsha256_BYTES],
                               uint8_t (*p_commitments)[][crypto_core_ristretto255_BYTES],
                               uint8_t (*kc0_commitments)[][crypto_core_ristretto255_BYTES],
                               uint8_t (*k0p_commitments)[][crypto_core_ristretto255_BYTES],
                               uint8_t (*zk_challenge_commitments)[][3][crypto_scalarmult_ristretto255_SCALARBYTES],
                               uint8_t (*zk_challenge_e_i)[][crypto_scalarmult_ristretto255_SCALARBYTES],
                               uint8_t (*k0p_final_commitments)[][crypto_scalarmult_ristretto255_BYTES],
                               uint64_t *last_ts);

/**
 * @enum TOPRF_Update_Peer_Steps
 * @brief Protocol steps for the peers
 *
 * These values define the sequential steps that each peer follows during
 * the execution of the TOPRF update protocol.
 */
typedef enum {
  TOPRF_Update_Peer_Broadcast_NPK_SIDNonce,
  TOPRF_Update_Peer_Rcv_NPK_SIDNonce,
  TOPRF_Update_Peer_Noise_Handshake,
  TOPRF_Update_Peer_Finish_Noise_Handshake,
  TOPRF_Update_Peer_Rcv_CHashes_Send_Commitments,
  TOPRF_Update_Peer_Rcv_Commitments_Send_Shares,
  TOPRF_Update_Peer_Verify_Commitments,
  TOPRF_Update_Peer_Handle_DKG_Complaints,
  TOPRF_Update_Peer_Defend_DKG_Accusations,
  TOPRF_Update_Peer_Check_Shares,
  TOPRF_Update_Peer_Finish_DKG,
  TOPRF_Update_Peer_Confirm_Transcripts,
  TOPRF_Update_Peer_Rcv_Mult_CHashes_Send_Commitments,
  TOPRF_Update_Peer_Send_K0P_Shares,
  TOPRF_Update_Peer_Recv_K0P_Shares,
  TOPRF_Update_Peer_Handle_Mult_Share_Complaints,
  TOPRF_Update_Peer_Defend_Mult_Accusations,
  TOPRF_Update_Peer_Check_Mult_Shares,
  TOPRF_Update_Peer_Disclose_Mult_Shares,
  TOPRF_Update_Peer_Reconstruct_Mult_Shares,
  TOPRF_Update_Peer_Send_ZK_Challenge_Commitments,
  TOPRF_Update_Peer_Send_ZK_Commitments,
  TOPRF_Update_Peer_Send_ZK_nonces,
  TOPRF_Update_Peer_Send_ZK_proofs,
  TOPRF_Update_Peer_Verify_ZK_proofs,
  TOPRF_Update_Peer_Disclose_ZK_Cheaters,
  TOPRF_Update_Peer_Reconstruct_ZK_Shares,
  TOPRF_Update_Peer_Send_Mult_Ci,
  TOPRF_Update_Peer_Final_VSPS_Checks,
  TOPRF_Update_Peer_Disclose_VSPS_Cheaters,
  TOPRF_Update_Peer_Reconstruct_VSPS_Shares,
  TOPRF_Update_Peer_Send_k0p_Share,
  TOPRF_Update_Peer_Final_OK,
  TOPRF_Update_Peer_Done
} TOPRF_Update_Peer_Steps;

/**
 * @enum TOPRF_Update_Message_Type
 * @brief Message type identifiers used in the TOPRF-Update protocol
 *
 * These identifiers are used to tag messages exchanged during the protocol,
 * allowing recipients to properly interpret and process each message.
 */
typedef enum {
  toprfupdate_stp_start_msg,
  toprfupdate_peer_init_msg,
  toprfupdate_stp_bc_init_msg,
  toprfupdate_peer_ake1_msg,
  toprfupdate_peer_ake2_msg,
  toprfupdate_peer_dkg1_msg,
  toprfupdate_stp_bc_dkg1_msg,
  toprfupdate_peer_dkg2_msg,
  toprfupdate_stp_bc_dkg2_msg,
  toprfupdate_peer_dkg3_msg,
  toprfupdate_peer_verify_shares_msg,
  toprfupdate_stp_bc_verify_shares_msg,
  toprfupdate_peer_share_key_msg,
  toprfupdate_stp_bc_key_msg,
  toprfupdate_peer_bc_transcript_msg,
  toprfupdate_stp_bc_transcript_msg,
  toprfupdate_peer_mult1_msg,
  toprfupdate_stp_bc_mult1_msg,
  toprfupdate_peer_mult_coms_msg,
  toprfupdate_stp_bc_mult_coms_msg,
  toprfupdate_peer_mult2_msg,
  toprfupdate_peer_verify_mult_shares_msg,
  toprfupdate_peer_share_mult_key_msg,
  toprfupdate_stp_bc_mult_key_msg,
  toprfupdate_peer_reconst_mult_shares_msg,
  toprfupdate_stp_bc_reconst_mult_shares_msg,
  toprfupdate_peer_zkp1_msg,
  toprfupdate_stp_bc_zkp1_msg,
  toprfupdate_peer_zkp2_msg,
  toprfupdate_stp_bc_zkp2_msg,
  toprfupdate_peer_zkp3_msg,
  toprfupdate_stp_bc_zkp3_msg,
  toprfupdate_peer_zkp4_msg,
  toprfupdate_stp_bc_zkp4_msg,
  toprfupdate_peer_zk_disclose_msg,
  toprfupdate_stp_bc_zk_disclose_msg,
  toprfupdate_peer_mult3_msg,
  toprfupdate_stp_bc_mult3_msg,
  toprfupdate_peer_vsps_disclose_msg,
  toprfupdate_stp_bc_vsps_disclose_msg,
  toprfupdate_peer_end2_msg,
  toprfupdate_stp_end3_msg,
} TOPRF_Update_Message_Type;

/**
 * @struct TOPRF_Update_ZK_params
 * @brief Parameters for zero-knowledge proofs
 *
 * This structure contains the parameters used in zero-knowledge proofs
 * within the TOPRF update protocol.
 */
typedef struct {
  uint8_t d[crypto_scalarmult_ristretto255_SCALARBYTES];
  uint8_t s[crypto_scalarmult_ristretto255_SCALARBYTES];
  uint8_t x[crypto_scalarmult_ristretto255_SCALARBYTES];
  uint8_t s_1[crypto_scalarmult_ristretto255_SCALARBYTES];
  uint8_t s_2[crypto_scalarmult_ristretto255_SCALARBYTES];
} TOPRF_Update_ZK_params;

/**
 * @struct TOPRF_Update_ZK_proof
 * @brief Values in a zero-knowledge proof
 *
 * This structure contains the public values used in a zero-knowledge
 * proof within the TOPRF update protocol.
 */
typedef struct {
  uint8_t y[crypto_scalarmult_ristretto255_SCALARBYTES];
  uint8_t w[crypto_scalarmult_ristretto255_SCALARBYTES];
  uint8_t z[crypto_scalarmult_ristretto255_SCALARBYTES];
  uint8_t w_1[crypto_scalarmult_ristretto255_SCALARBYTES];
  uint8_t w_2[crypto_scalarmult_ristretto255_SCALARBYTES];
} TOPRF_Update_ZK_proof;

/**
 * @struct TOPRF_Update_PeerState
 * @brief State of a peer during the TOPRF update protocol
 *
 * This contains the state of a peer during the execution of the TOPRF
 * Update protocol.
 *
 * Some fields in this struct are internal variables and should not
 * be used. The following fields are useful and can be accessed by
 * users of the API:
 *
 * @var TOPRF_Update_PeerState::n Total number of peers participating in
 *      this protocol
 *
 * @var TOPRF_Update_PeerState::t Threshold necessary to use the shared
 *      secret generated by this DKG
 *
 * @var TOPRF_Update_PeerState::index The index of the peer
 *
 * @var TOPRF_Update_PeerState::cheaters List of detected cheaters and
 *      protocol violators at the end of the failed protocol run
 */
typedef struct {
  TOPRF_Update_Peer_Steps step;
  TOPRF_Update_Peer_Steps prev;
  uint8_t sessionid[dkg_sessionid_SIZE];
  uint8_t n;
  uint8_t t;
  uint8_t index;
  TOPRF_Share kc0_share[2];
  uint8_t (*kc0_commitments)[][crypto_core_ristretto255_BYTES];
  uint8_t sig_pk[crypto_sign_PUBLICKEYBYTES];
  uint8_t sig_sk[crypto_sign_SECRETKEYBYTES];
  uint8_t noise_pk[crypto_scalarmult_BYTES];
  uint8_t noise_sk[crypto_scalarmult_SCALARBYTES];
  uint64_t stp_last_ts;
  uint64_t *last_ts;
  uint64_t ts_epsilon;
  const uint8_t (*sig_pks)[][crypto_sign_PUBLICKEYBYTES];
  uint8_t (*peer_noise_pks)[][crypto_scalarmult_BYTES];
  Noise_XK_device_t *dev;
  Noise_XK_session_t *(*noise_outs)[];
  Noise_XK_session_t *(*noise_ins)[];

  uint8_t (*encrypted_shares)[][noise_xk_handshake3_SIZE + toprf_update_encrypted_shares_SIZE];

  TOPRF_Share (*p_shares)[][2];
  uint8_t (*p_commitments)[][crypto_core_ristretto255_BYTES];
  uint8_t (*p_commitments_hashes)[][toprf_update_commitment_HASHBYTES];
  uint8_t (*p_share_macs)[][crypto_auth_hmacsha256_BYTES];
  uint16_t p_complaints_len;
  uint16_t *p_complaints;
  uint8_t my_p_complaints_len;
  uint8_t *my_p_complaints;
  TOPRF_Share p_share[2];

  uint8_t (*lambdas)[][crypto_core_ristretto255_SCALARBYTES];
  TOPRF_Share (*k0p_shares)[][2];
  uint8_t (*k0p_commitments)[][crypto_core_ristretto255_BYTES];
  uint8_t k0p_tau[crypto_core_ristretto255_SCALARBYTES];

  uint8_t zk_chal_nonce[2][crypto_core_ristretto255_SCALARBYTES];
  uint8_t (*zk_challenge_nonces)[][2][crypto_scalarmult_ristretto255_SCALARBYTES];
  uint8_t (*zk_challenge_nonce_commitments)[][crypto_scalarmult_ristretto255_BYTES];
  uint8_t (*zk_challenge_commitments)[][3][crypto_scalarmult_ristretto255_SCALARBYTES];
  uint8_t (*zk_challenge_e_i)[][crypto_scalarmult_ristretto255_SCALARBYTES];
  TOPRF_Update_ZK_params zk_params;

  TOPRF_Share k0p_share[2];
  uint8_t k0p_commitment[crypto_core_ristretto255_BYTES];

  size_t cheater_len;
  TOPRF_Update_Cheater (*cheaters)[];
  size_t cheater_max;
  crypto_generichash_state transcript_state;
  uint8_t transcript[crypto_generichash_BYTES];
} TOPRF_Update_PeerState;

/**
 * @brief Gets the size of a TOPRF_Update_PeerState struct
 *
 * Returns the memory size required for the PeerState structure.
 * Useful when dynamically allocating memory for this structure. Note that
 * the struct must be aligned at a 32-byte boundary for proper operation.
 *
 * @return Size of the TOPRF_Update_PeerState structure in bytes
 */
size_t toprf_update_peerstate_size(void);

/**
 * @brief Gets the value of `n` (number of participants) from the peer state
 *
 * @param ctx Pointer to an initialized TOPRF_Update_PeerState structure
 *
 * @return The number of participants in the protocol
 */
uint8_t toprf_update_peerstate_n(const TOPRF_Update_PeerState *ctx);

/**
 * @brief Gets the value of T (threshold) from the peer state
 *
 * @param ctx Pointer to an initialized TOPRF_Update_PeerState structure
 *
 * @return The threshold value used in the protocol
 */
uint8_t toprf_update_peerstate_t(const TOPRF_Update_PeerState *ctx);

/**
 * @brief Gets the session ID from the peer state
 *
 * @param ctx Pointer to an initialized TOPRF_Update_PeerState structure
 *
 * @return Pointer to the session ID byte array
 */
const uint8_t *toprf_update_peerstate_sessionid(const TOPRF_Update_PeerState *ctx);

/**
 * @brief Gets the private share of the updated OPRF key
 *
 * Returns the final result, the peer's private share of the updated
 * OPRF key from a TOPRF_Update_PeerState struct. This should only be
 * called after the protocol has finished successfully.
 *
 * @param ctx Pointer to an initialized TOPRF_Update_PeerState structure
 * @return Pointer to the share byte array, or NULL if protocol not complete
 */
const uint8_t *toprf_update_peerstate_share(const TOPRF_Update_PeerState *ctx);

/**
 * @brief Gets the commitment to the peer's share of the updated OPRF key
 *
 * Returns the public commitment to the peer's share of the updated
 * OPRF key from a TOPRF_Update_PeerState struct. This should only be
 * called after the protocol has finished successfully.
 *
 * @param ctx Pointer to an initialized TOPRF_Update_PeerState structure
 *
 * @return Pointer to the commitment byte array, or NULL if protocol
 *         not complete
 */
const uint8_t *toprf_update_peerstate_commitment(const TOPRF_Update_PeerState *ctx);

/**
 * @brief Gets all the public commitments to the shares of the updated
 *        OPRF key
 *
 * Returns the public commitments to all shares of the updated OPRF key
 * from a TOPRF_Update_PeerState struct. This should only be used after
 * the protocol has finished successfully.
 *
 * @param ctx Pointer to an initialized TOPRF_Update_PeerState structure
 * @return Pointer to the commitments byte array, or NULL if protocol not complete
 */
const uint8_t *toprf_update_peerstate_commitments(const TOPRF_Update_PeerState *ctx);

/**
 * @brief Gets the current step of the protocol from the peer state
 *
 * @param ctx Pointer to an initialized TOPRF_Update_PeerState structure
 *
 * @return The current protocol step as an integer
 */
int toprf_update_peerstate_step(const TOPRF_Update_PeerState *ctx);

/**
 * @brief Starts a new execution of the TOPRF-Update protocol for a peer
 *
 * This function initializes the state of the peer.
 *
 * @param[in] ctx Pointer to a TOPRF_Update_PeerState struct. This
              struct will be initialized by this function
 * @param[in] ts_epsilon Maximum allowed message age in seconds before
 *            it is considered stale and rejected. This value is used to
 *            prevent replay attacks and enforce freshness. For small,
 *            local setups (e.g., 2-out-of-3 participants), values as low
 *            as 2–3 seconds may suffice. For large-scale deployments
 *            (e.g., 126-out-of-127), this may need to be increased to
 *            several hours
 * @param[in] lt_sk Long-term private signing key of the peer
 * @param[in] noise_sk The Noise_XK private key of the peer
 * @param[in] msg0 The initiating message received from the TP (created
 *            after running `tpdkg_tp_start()`)
 * @param[out] keyid The ID of the key to be updated. Use this to look up
 *             the `n-t` parameters, the key to be updated, and the various
 *             long-term public keys of the other participants
 * @param[out] stp_ltpk The STP's long-term public signing key. This should
 *            be used to verify if this key is actually authorized to
 *            initiate an STP DKG with the peer
 *
 * @return 0 on success, non-zero on error
 */
TOPRF_Update_Err toprf_update_start_peer(TOPRF_Update_PeerState *ctx,
                                         const uint64_t ts_epsilon,
                                         const uint8_t lt_sk[crypto_sign_SECRETKEYBYTES],
                                         const uint8_t noise_sk[crypto_scalarmult_SCALARBYTES],
                                         const TOPRF_Update_Message *msg0,
                                         uint8_t keyid[toprf_keyid_SIZE],
                                         uint8_t stp_ltpk[crypto_sign_PUBLICKEYBYTES]);

/**
 * @brief Sets all the variable sized buffers in the TOPRF_Update_PeerState
 *        structure
 *
 * This function sets up pointers to all variable-sized buffers required by
 * a peer during the TOPRF key update protocol.
 *
 * The size of most buffers depends on the `n` (number of participants) and
 * `t` (threshold) values, which are derived from the DKG parameters and
 * communicated in `msg0` from the STP. These are inputs to
 * `toprf_update_start_peer()`, and must be passed to this function as
 * parameters.
 * After calling this function, the `TOPRF_Update_PeerState` struct is ready
 * to be used, and the caller can query the parameters `n`, `t` and the
 * `keyid` to load the corresponding data.
 *
 * To allocate all the buffers on the stack:
 * @code
 * Noise_XK_session_t *noise_outs[n];
 * Noise_XK_session_t *noise_ins[n];
 *
 * TOPRF_Share pshares[n][2];
 * uint8_t p_commitments[n*n][crypto_core_ristretto255_BYTES];
 * uint8_t p_commitments_hashes[n][toprf_update_commitment_HASHBYTES];
 * uint8_t peers_p_share_macs[n*n][crypto_auth_hmacsha256_BYTES];
 * uint16_t peer_p_complaints[n*n];
 * uint8_t peer_my_p_complaints[n];
 *
 * uint8_t encrypted_shares[n][noise_xk_handshake3_SIZE + toprf_update_encrypted_shares_SIZE];
 *
 * uint64_t peer_last_ts[n];
 * uint8_t lambdas[dealers][crypto_core_ristretto255_SCALARBYTES];
 * TOPRF_Share k0p_shares[dealers][2];
 * uint8_t k0p_commitments[dealers*(n+1)][crypto_core_ristretto255_BYTES];
 * uint8_t zk_challenge_nonce_commitments[n*2][crypto_scalarmult_ristretto255_BYTES];
 * uint8_t zk_challenge_nonces[n*2][2][crypto_scalarmult_ristretto255_SCALARBYTES];
 * uint8_t zk_challenge_commitments[dealers*2][3][crypto_scalarmult_ristretto255_SCALARBYTES];
 * uint8_t zk_challenge_e_i[2*dealers][crypto_scalarmult_ristretto255_SCALARBYTES];
 * TOPRF_Update_Cheater peer_cheaters[n*n - 1];
 *
 * if(0!=toprf_update_peer_set_bufs(&peer, index, n, t, k0_share,
 *                                  &k0_commitments,
 *                                   &lt_pks, &peers_noise_pks, noise_sk,
 *                                  &noise_outs, &noise_ins,
 *                                  &pshares,
 *                                  &p_commitments,
 *                                  &p_commitments_hashes,
 *                                  &peers_p_share_macs,
 *                                  &encrypted_shares,
 *                                  &peer_cheaters, sizeof(peer_cheaters) / sizeof(TOPRF_Update_Cheater) / n,
 *                                  &lambdas,
 *                                  &k0p_shares,
 *                                  &k0p_commitments,
 *                                  &zk_challenge_nonce_commitments,
 *                                  &zk_challenge_nonces,
 *                                  &zk_challenge_commitments,
 *                                  &zk_challenge_e_i,
 *                                  peer_p_complaints,
 *                                  peer_my_p_complaints,
 *                                  peer_last_ts)) return 1;
 * @endcode
 *
 * @param[in] index Must be initialized based on the index of the `k0`
 *             share index. If no share is held by the peer, the index
 *             must be established out of band. TODO: fix this
 * @param[in] n The total number of participants in the protocol. It must
 *           be initialized based on the info associated with the `keyid`
 *           returned by `toprf_update_start_peer()`
 * @param[in] t Threshold of participants required to reconstruct the
 *            secret. It must be initialized based on the info associated
 *            with the `keyid` returned by `toprf_update_start_peer()`
 * @param[in] k0 Pointer to the old key shares held by this peer). Pass
 *            NULL if the peer holds no shares.
 * @param[in] kc0_commitments The commitments to the shares of the old key.
 *            Even if peer has no share, they must have these commitments.
 *            TODO: check if this is actually true.
 * @param[in] sig_pks Array of long-term signing public keys for all
 *            participants. The STP's public key must be at index 0. All
 *            other public keys must be in order of the shares of `k0` they
 *            hold. Peers that hold no `k0` shares can be listed in
 *            arbitrary order, following those with shares
 * @param[in] peer_noise_pks Array of all Noise_XK public keys of peers,
 *            in the order of the indices of the shares of `k0`. Peers that
 *            hold no `k0` shares can be listed in arbitrary order, following
 *            those with shares.
 *
 * @return 0 on success, non-zero on error
 *
 */
int toprf_update_peer_set_bufs(TOPRF_Update_PeerState *ctx,
                               const uint8_t self,
                               const uint8_t n, const uint8_t t,
                               const TOPRF_Share k0[2],
                               uint8_t (*k0_commitments)[][crypto_core_ristretto255_BYTES],
                               const uint8_t (*sig_pks)[][crypto_sign_PUBLICKEYBYTES],
                               uint8_t (*noise_pks)[][crypto_scalarmult_BYTES],
                               Noise_XK_session_t *(*noise_outs)[],
                               Noise_XK_session_t *(*noise_ins)[],
                               TOPRF_Share (*p_shares)[][2],
                               uint8_t (*p_commitments)[][crypto_core_ristretto255_BYTES],
                               uint8_t (*p_commitments_hashes)[][toprf_update_commitment_HASHBYTES],
                               uint8_t (*p_share_macs)[][crypto_auth_hmacsha256_BYTES],
                               uint8_t (*encrypted_shares)[][noise_xk_handshake3_SIZE + toprf_update_encrypted_shares_SIZE],
                               TOPRF_Update_Cheater (*cheaters)[], const size_t cheater_max,
                               uint8_t (*lambdas)[][crypto_core_ristretto255_SCALARBYTES],
                               TOPRF_Share (*k0p_shares)[][2],
                               uint8_t (*k0p_commitments)[][crypto_core_ristretto255_BYTES],
                               uint8_t (*zk_challenge_nonce_commitments)[][crypto_scalarmult_ristretto255_BYTES],
                               uint8_t (*zk_challenge_nonces)[][2][crypto_scalarmult_ristretto255_SCALARBYTES],
                               uint8_t (*zk_challenge_commitments)[][3][crypto_scalarmult_ristretto255_SCALARBYTES],
                               uint8_t (*zk_challenge_e_i)[][crypto_scalarmult_ristretto255_SCALARBYTES],
                               uint16_t *p_complaints,
                               uint8_t *my_p_complaints,
                               uint64_t *last_ts);
/**
 * @brief Calculates the size of the buffer needed to hold all outputs
 *        from peers, which serve as input to the next step of the STP
 *
 * An implementer should allocate a buffer of this size and concatenate
 * all messages from all peers in the order of the peers. The allocated
 * buffer is to be passed as an input to the `toprf_update_stp_next()`
 * function. After this, the buffer SHOULD be deallocated.
 *
 * @param[in] ctx An initialized TOPRF_Update_STPState struct
 *
 * @return 1 on error, otherwise the size to be allocated (can be 0)
 */
size_t toprf_update_stp_input_size(const TOPRF_Update_STPState *ctx);

/**
 * @brief Calculates the size of the message from each peer to be received
 *        by the STP
 *
 * This function fills the `sizes` array (which must have exactly `n`
 * elements, one per peer) with the expected message size from each
 * peer based on the current STP step.
 *
 * @param[in]  ctx   An initialized TOPRF_Update_STPState struct
 * @param[out] sizes An array of type `size_t` with exactly `n` elements
 *
 * @return 0 on if the sizes differ from peer to peer, otherwise all
 *         peers will be sending messages of equal size. In the latter
 *         case, all items of the sizes array hold the same valid value
 */
int toprf_update_stp_input_sizes(const TOPRF_Update_STPState *ctx, size_t *sizes);

/**
 * @brief Calculates the size of the buffer needed to hold the output
 *        from the `toprf_update_stp_next()` function
 *
 * An implementer should allocate a buffer of this size and pass if as
 * parameter to `toprf_update_stp_next()`.
 *
 * @param[in] ctx Pointer to an initialized TOPRF_Update_STPState struct
 *
 * @return 1 on error, otherwise the size to be allocated (can be 0)
 */
size_t toprf_update_stp_output_size(const TOPRF_Update_STPState *ctx);

/**
 * @brief Executes the next step of the TOPRF update protocol for the STP
 *
 * This function handles one round of the protocol on the STP side by
 * reading peer input, processing it, and writing output that may need
 * to be sent to peers.
 *
 * @param[in] ctx Pointer to an initialized TOPRF_Update_STPState struct
 * @param[in] input Pointer to input buffer for this step
 * @param[in] input_len Size of the input buffer
 * @param[out] output Pointer to the output buffer for this step
 * @param[in] output_len Size of the output buffer
 *
 * @return 0 on success, non-zero on error
 *
 * An example of how to use this function in concert with
 * `toprf_update_stp_input_size()` and `toprf_update_stp_output_size()`:
 * @code
 * uint8_t stp_out[toprf_update_stp_output_size(&tp)];
 * uint8_t stp_in[toprf_update_stp_input_size(&tp)];
 * recv(socket, stp_in, sizeof(stp_in));
 * ret = toprf_update_stp_next(&stp, stp_in, sizeof(stp_in),
 *                              stp_out, sizeof(stp_out));
 * @endcode
 */
int toprf_update_stp_next(TOPRF_Update_STPState *ctx, const uint8_t *input, const size_t input_len, uint8_t *output, const size_t output_len);

/**
 * @brief Extracts the per-peer message from the output of
 *        `toprf_update_stp_next()` for the specified peer
 *
 * This function converts the output of `toprf_update_stp_next()` into a
 * message to be sent to the i-th peer. Some outputs of steps produce
 * broadcast messages (same messages for all peers), while others generate
 * decicated and unique messages for each peer.
 * This function returns a pointer to a message and the size of the message
 * to be sent for a particular peer specified as a parameter.
 *
 * Example of how to use this function in concert with
`toprf_update_stp_next()`:
 * @code
 *  ret = toprf_update_stp_next(&stp, stp_in, sizeof(stp_in), stp_out, sizeof stp_out);
 *
 *  for(int i=0;i<stp.n;i++) {
 *    const uint8_t *msg;
 *    size_t len;
 *    if(0!=toprf_update_stp_peer_msg(&stp, stp_out, sizeof stp_out, i, &msg, &len)) {
 *      return 1;
 *    }
 *    send(i, msg, len);
 *  }
 * @endcode
 *
 * @param[in] ctx Pointer to a valid TOPRF_Update_STPState struct
 * @param[in] base Pointer to the output of `toprf_update_stp_next()`
 * @param[in] base_size Size (in bytes) of the output of
 *             `toprf_update_stp_next()`
 * @param[in] peer Index of the peer (0-based)
 * @param[out] msg Pointer to the message to be sent to the i-th peer
 * @param[out] len Pointer to the length of the message to be sent to the
 *             i-th peer
 *
 * @return 0 on success, non-zero on error
 *
 */
int toprf_update_stp_peer_msg(const TOPRF_Update_STPState *ctx, const uint8_t *base, const size_t base_size, const uint8_t peer, const uint8_t **msg, size_t *len);

/**
 * @brief Checks if the protocol has finished for the STP or more
 *        `toprf_update_stp_next()` calls are necessary
 *
 * @param stp Pointer to the STP state structure.
 *
 * @return 1 if more steps are outstanding
 */
int toprf_update_stp_not_done(const TOPRF_Update_STPState *stp);

/**
 * @brief Calculates the size of the buffer needed to hold the output
 *        from the STP serving as input to the next step of the peer
 *
 * An implementer should allocate a buffer of this size. The allocated
 * buffer is to be passed as an input to the `toprf_update_peer_next()`
 * function. After this, the buffer SHOULD be deallocated.
 *
 * @param[in] ctx Pointer to an initialized TOPRF_Update_PeerState struct
 *
 * @return 1 on error, otherwise the size to be allocated (can be 0)
 */
size_t toprf_update_peer_input_size(const TOPRF_Update_PeerState *ctx);

/**
 * @brief Calculates the size of the buffer needed to hold the output
 *        from the `toprf_update_peer_next()` function
 *
 * An implementer should allocate a buffer of this size and pass it as parameter to
 * `toprf_update_peer_next()`
 *
 * @param[in] ctx Pointer to initialized TOPRF_Update_PeerState struct
 *
 * @return 1 on error, otherwise the size to be allocated (can be 0)
 */
size_t toprf_update_peer_output_size(const TOPRF_Update_PeerState *ctx);

/**
 * @brief Executes the next step in the TOPRF Update protocol for a peer
 *
 * On success, the protocol advances to the next step automatically.
 * On error, the peer state is moved to an invalid state to mark the
 * failure.
 *
 * An example of how to use this function in concert with
 * `toprf_update_peer_input_size()` and `toprf_update_peer_output_size()`
 * while allocating the buffers on the stack:
 * @code
 * uint8_t peers_out[toprf_update_peer_output_size(&peer)];
 * uint8_t peer_in[toprf_update_peer_input_size(&peer)];
 * recv(socket, peer_in, sizeof(peer_in));
 * int ret = toprf_update_peer_next(&peer,
 *                                   peer_in, sizeof(peer_in),
 *                                   peers_out, sizeof(peers_out));
 * @endcode
 *
 * @param[in] ctx Pointer to initialized TOPRF_Update_PeerState
 * @param[in] input Pointer to input buffer for the current step
 * @param[in] input_len Length of input buffer
 * @param[out] output Pointer to output buffer for the current step
 * @param[in] output_len Length of output buffer
 *
 * @return 0 on success, non-zero otherwise
 */
int toprf_update_peer_next(TOPRF_Update_PeerState *ctx, const uint8_t *input, const size_t input_len, uint8_t *output, const size_t output_len);

/**
 * @brief Checks if the protocol has finished for the peer or more
 * `toprf_update_peer_next()` calls are necessary
 *
 * @param peer Pointer to the peer state structure
 *
 * @return 1 if the peer update step is not done, 0 otherwise
 */
int toprf_update_peer_not_done(const TOPRF_Update_PeerState *peer);

/**
 * @brief Frees dynamically allocated Noise_XK sessions associated with a peer context
 *
 * This function MUST be used before a peer's state is deallocated.
 * The underlying Noise_XK implementation allocates internal state on the heap,
 * which must be freed manually to avoid memory leaks.
 *
 * @param ctx Pointer to the peer state structure
 */
void toprf_update_peer_free(TOPRF_Update_PeerState *ctx);

/**
   log_file is a FILE* pointing to an open file or NULL.

   If the file is not NULL the implementation will emit log messages
   to this file.
 */
extern FILE* log_file;

#endif //TOPRF_UPDATE_H
