use miden::core::crypto::dsa::ecdsa_k256_keccak
use miden::core::crypto::hashes::rpo256
use miden::protocol::active_account
use miden::protocol::native_account
use miden::protocol::tx
use miden::standards::auth

# CONSTANTS
# =================================================================================================

# The event to request an authentication signature.
const AUTH_REQUEST_EVENT=event("miden::auth::request")

# Local Memory Addresses for multisig operations
const NUM_OF_APPROVERS_LOC=0
const PUB_KEY_SLOT_SUFFIX_LOC=4
const PUB_KEY_SLOT_PREFIX_LOC=5
const CURRENT_PK_LOC=8
const SUCCESSFUL_VERIFICATIONS_LOC=12

#! Authenticate a transaction using the ECDSA signature scheme.
#!
#! It first increments the nonce of the account, independent of whether the account's state has
#! changed or not. Then it computes and signs the following message (in memory order):
#! [ACCOUNT_DELTA_COMMITMENT, INPUT_NOTES_COMMITMENT,
#!  OUTPUT_NOTES_COMMITMENT, [0, 0, ref_block_num, final_nonce]]
#!
#! Including the final_nonce is necessary for replay protection. The reference block number is
#! included to commit to the transaction creator's intended reference block of the transaction
#! which determines the fee parameters and therefore the fee amount that is deducted.
#!
#! Inputs:  [PUB_KEY]
#! Outputs: []
#!
#! Invocation: exec
pub proc authenticate_transaction
    # Increment the account's nonce.
    # ---------------------------------------------------------------------------------------------
    # This has to happen before computing the delta commitment, otherwise that procedure will abort
    push.0.0 exec.tx::get_block_number
    exec.native_account::incr_nonce
    # => [[final_nonce, ref_block_num, 0, 0], PUB_KEY]

    # Compute the message that is signed.
    # ---------------------------------------------------------------------------------------------
    exec.auth::create_tx_summary
    # => [SALT, OUTPUT_NOTES_COMMITMENT, INPUT_NOTES_COMMITMENT, ACCOUNT_DELTA_COMMITMENT, PUB_KEY]

    exec.auth::adv_insert_hqword

    # The commitment to the tx summary is the message that is signed
    exec.auth::hash_tx_summary
    # OS => [MESSAGE, PUB_KEY]
    # AS => []

    # Fetch signature from advice provider and verify.
    # ---------------------------------------------------------------------------------------------
    # Emit the authentication request event that pushes a signature for the message to the advice stack
    emit.AUTH_REQUEST_EVENT
    swapw
    # OS => [PUB_KEY, MESSAGE]
    # AS => [SIGNATURE]

    # Verify the signature against the public key and the message. The procedure gets as inputs the
    # hash of the public key and the message via the operand stack. The signature is provided via
    # the advice stack. The signature is valid if and only if the procedure returns.
    exec.ecdsa_k256_keccak::verify
    # OS => []
    # AS => []
end

#! Verify signatures for all required signers in a loop.
#!
#! This procedure iterates through the required number of signers, fetches their public keys from
#! the provided account storage map slot, verifies their signatures against the transaction message,
#! and returns the number of successfully verified signatures.
#!
#! Note: Calls `active_account::get_initial_map_item` to access the transaction's initial storage
#! state rather than the current state. This is crucial when validating transactions that update
#! the owner public key mapping - the previous signers must authorize the change to the new signers,
#! not the new signers authorizing themselves.
#!
#! Inputs:  [pub_key_slot_prefix, pub_key_slot_suffix, num_of_approvers, MSG]
#! Outputs: [num_verified_signatures, MSG]
@locals(16)
pub proc verify_signatures
    loc_store.PUB_KEY_SLOT_PREFIX_LOC
    loc_store.PUB_KEY_SLOT_SUFFIX_LOC
    # => [num_of_approvers, MSG]

    # Initializing SUCCESSFUL_VERIFICATIONS local memory address to 0
    push.0 loc_store.SUCCESSFUL_VERIFICATIONS_LOC
    # => [num_of_approvers, MSG]

    # Counter `i` starts at `num_of_approvers` and counts down to 0
    # => [i, MSG]

    # Loop through required signers and verify signatures.
    # ---------------------------------------------------------------------------------------------

    dup neq.0
    while.true
        # => [i, MSG]

        # Fetch public key from storage map.
        # -----------------------------------------------------------------------------------------

        sub.1 dup push.0.0.0
        loc_load.PUB_KEY_SLOT_SUFFIX_LOC loc_load.PUB_KEY_SLOT_PREFIX_LOC
        # => [owner_key_slot_prefix, owner_key_slot_suffix, [0, 0, 0, i-1], i-1, MSG]

        # Get public key from initial storage state
        exec.active_account::get_initial_map_item
        # => [OWNER_PUB_KEY, i-1, MSG]

        loc_storew_be.CURRENT_PK_LOC
        # => [OWNER_PUB_KEY, i-1, MSG]

        # Check if signature exists for this signer.
        # -----------------------------------------------------------------------------------------

        movup.4 movdn.8 swapw dupw movdnw.2
        # => [MSG, OWNER_PUB_KEY, MSG, i-1]

        exec.rpo256::merge
        # => [SIG_KEY, MSG, i-1]

        adv.has_mapkey
        # => [SIG_KEY, MSG, i-1]

        adv_push.1
        # => [has_signature, SIG_KEY, MSG, i-1]

        # if SIG_KEY => SIGNATURE exists in AdviceMap check the signature
        if.true
            # => [SIG_KEY, MSG, i-1]

            dupw.1
            # => [MSG, SIG_KEY, MSG, i-1]

            swapw
            # => [SIG_KEY, MSG, MSG, i-1]

            # Verify the signature against the public key and message.
            # -----------------------------------------------------------------------------------------

            loc_loadw_be.CURRENT_PK_LOC
            # => [PK, MSG, MSG, i-1]

            swapw
            # => [MSG, PK, MSG, i-1]

            # Emit the authentication request event that pushes a signature for the message to the advice stack.
            emit.AUTH_REQUEST_EVENT

            swapw
            # OS => [PUB_KEY, MSG, MSG, i-1]
            # AS => [SIGNATURE]

            exec.ecdsa_k256_keccak::verify
            # => [MSG, i-1]

            loc_load.SUCCESSFUL_VERIFICATIONS_LOC
            add.1
            loc_store.SUCCESSFUL_VERIFICATIONS_LOC
            # => [MSG, i-1]
        else
            dropw
            # => [MSG, i-1]
        end
        # => [MSG, i-1]

        movup.4
        # => [i-1, MSG]

        dup neq.0
        # => [should_continue, i-1, MSG]
    end
    # => [i-1, MSG]

    # Return successful signature verifications along with MSG
    # -----------------------------------------------------------------------------------------

    drop
    # => [MSG]

    loc_load.SUCCESSFUL_VERIFICATIONS_LOC
    # => [num_verified_signatures, MSG]
end
