use.kernel::account
use.kernel::asset
use.kernel::constants
use.kernel::memory
use.kernel::note

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

# Constants for different note types
const.PUBLIC_NOTE=1     # 0b01
const.PRIVATE_NOTE=2    # 0b10
const.ENCRYPTED_NOTE=3  # 0b11

# Two raised to the power of 38 (2^38), used for shifting the note type value
const.TWO_POW_38=274877906944

# Max value for U16, used as the upper limit for expiration block delta
const.EXPIRY_UPPER_LIMIT=0xFFFF+1

# The note type must be PUBLIC, unless the high bits are `0b11`. (See the table below.)
const.ALL_NOTE_TYPES_ALLOWED=3 # 0b11

# Max U32 value, used for initializing the expiration block number
const.MAX_BLOCK_NUM=0xFFFFFFFF

# ERRORS
# =================================================================================================

# Number of output notes in the transaction exceeds the maximum limit of 1024
const.ERR_TX_NUMBER_OF_OUTPUT_NOTES_EXCEEDS_LIMIT=0x00020042

# Invalid note type
const.ERR_NOTE_INVALID_TYPE=0x00020043

# The 2 highest bits in the u32 tag have the following meaning:
#
# | bits | Execution | Target    | Note type |
# | ---- | --------- | --------- | --------- |
# | 00   | Network   | Specific  | Public    |
# | 01   | Network   | Use case  | Public    |
# | 10   | Local     | Any       | Public    |
# | 11   | Local     | Any       | Any       |
#
# Execution: Is a hint for the network, to check if the note can be consumed by a network controlled
#   account
# Target: Is a hint for the type of target. Use case means the note may be consumed by anyone,
#   specific means there is a specific target for the note (the target may be a public key, a user
#   that knows some secret, or a specific account id)
#
# Only the note type from the above list is enforced. The other values are only hints intended as a
# best effort optimization strategy. A badly formatted note may 1. not be consumed because honest
# users won't see the note 2. generate slightly more load as extra validation is performed for the
# invalid tags. None of these scenarios have any significant impact.

# Invalid note type for the given note tag prefix
const.ERR_NOTE_INVALID_NOTE_TYPE_FOR_NOTE_TAG_PREFIX=0x00020044

# The note's tag must fit into a u32 so the 32 most significant bits must be zero.
const.ERR_NOTE_TAG_MUST_BE_U32=0x00020045

# Adding a fungible asset to a note cannot exceed the max_amount of 9223372036854775807
const.ERR_NOTE_FUNGIBLE_MAX_AMOUNT_EXCEEDED=0x00020046

# Non-fungible asset that already exists in the note cannot be added again
const.ERR_NON_FUNGIBLE_ASSET_ALREADY_EXISTS=0x00020047

# Failed to find note at the given index; index must be within [0, num_of_notes]
const.ERR_NOTE_INVALID_INDEX=0x00020048

# Transaction expiration block delta must be within 0x1 and 0xFFFF.
const.ERR_TX_INVALID_EXPIRATION_DELTA=0x00020049

# EVENTS
# =================================================================================================

# Event emitted before a new note is created.
const.NOTE_BEFORE_CREATED_EVENT=131083
# Event emitted after a new note is created.
const.NOTE_AFTER_CREATED_EVENT=131084

# Event emitted before an ASSET is added to a note
const.NOTE_BEFORE_ADD_ASSET_EVENT=131085
# Event emitted after an ASSET is added to a note
const.NOTE_AFTER_ADD_ASSET_EVENT=131086

#! Returns the block hash of the reference block to memory.
#!
#! Stack: []
#! Output: [BLOCK_HASH]
#!
#! Where:
#! - BLOCK_HASH, reference block for the transaction execution.
export.memory::get_block_hash

#! Returns the block number of the last known block at the time of transaction execution.
#!
#! Inputs: []
#! Outputs: [num]
#!
#! num is the last known block number.
export.memory::get_blk_num->get_block_number

#! Returns the input notes commitment hash.
#!
#! See `transaction::api::get_input_notes_commitment` for details.
#!
#! Inputs: []
#! Outputs: [INPUT_NOTES_COMMITMENT]
#!
#! Where:
#! - INPUT_NOTES_COMMITMENT is the input notes commitment hash.
export.memory::get_input_notes_commitment

#! Returns the output notes hash. This is computed as a sequential hash of (note_id, note_metadata)
#! tuples over all output notes.
#!
#! Inputs: []
#! Outputs: [COM]
#!
#! COM is the output notes hash.
export.note::compute_output_notes_commitment->get_output_notes_hash

#! Increments the number of output notes by one. Returns the index of the next note to be created.
#!
#! Inputs: []
#! Outputs: [note_idx]
proc.increment_num_output_notes
    # get the current number of output notes
    exec.memory::get_num_output_notes
    # => [note_idx]

    # assert that there is space for a new note
    dup exec.constants::get_max_num_output_notes lt assert.err=ERR_TX_NUMBER_OF_OUTPUT_NOTES_EXCEEDS_LIMIT
    # => [note_idx]

    # increment the number of output notes
    dup add.1 exec.memory::set_num_output_notes
    # => [note_idx]
end

#! Adds a non-fungible asset to a note at the next available position.
#! Returns the pointer to the note the asset was stored at.
#! Panics if the non-fungible asset already exists in the note.
#!
#! Inputs: [ASSET, note_ptr, num_of_assets, note_idx]
#! Outputs: [note_ptr, note_idx]
proc.add_non_fungible_asset_to_note
    dup.4 exec.memory::get_output_note_asset_data_ptr
    # => [asset_ptr, ASSET, note_ptr, num_of_assets, note_idx]

    # compute the pointer at which we should stop iterating
    dup dup.7 add
    # => [end_asset_ptr, asset_ptr, ASSET, note_ptr, num_of_assets, note_idx]

    # reorganize and pad the stack, prepare for the loop
    movdn.5 movdn.5 padw dup.9
    # => [asset_ptr, 0, 0, 0, 0, ASSET, end_asset_ptr, asset_ptr, note_ptr, num_of_assets, note_idx]

    # compute the loop latch
    dup dup.10 neq
    # => [latch, asset_ptr, 0, 0, 0, 0, ASSET, end_asset_ptr, asset_ptr, note_ptr, num_of_assets, note_idx]

    while.true
        # load the asset and compare
        mem_loadw eqw assertz.err=ERR_NON_FUNGIBLE_ASSET_ALREADY_EXISTS
        # => [ASSET', ASSET, end_asset_ptr, asset_ptr, note_ptr, num_of_assets, note_idx]

        # drop ASSET' and increment the asset pointer
        dropw movup.5 add.1 dup movdn.6 padw movup.4
         # => [asset_ptr + 1, 0, 0, 0, 0, ASSET, end_asset_ptr, asset_ptr + 1, note_ptr, num_of_assets, note_idx]

        # check if we reached the end of the loop
        dup dup.10 neq
    end
    # => [asset_ptr, 0, 0, 0, 0, ASSET, end_asset_ptr, asset_ptr, note_ptr, num_of_assets, note_idx]

    # prepare stack for storing the ASSET
    movdn.4 dropw
    # => [asset_ptr, ASSET, end_asset_ptr, asset_ptr, note_ptr, num_of_assets, note_idx]

    # end of the loop reached, no error so we can store the non-fungible asset
    mem_storew dropw drop drop
    # => [note_ptr, num_of_assets, note_idx]

    # increase the number of assets in the note
    swap add.1 dup.1 exec.memory::set_output_note_num_assets
    # => [note_ptr, note_idx]
end

#! Updates the transaction expiration block number.
#!
#! The input block_height_delta is added to the block reference number in order to output an upper
#! limit at which the transaction will be considered valid (not expired).
#! This value can be later decreased, but not increased.
#!
#! Inputs: [block_height_delta, ...]
#! Output: [...]
#!
#! Where:
#! - block_height_delta is the desired expiration time delta (1 to 0xFFFF).
export.update_expiration_block_num
    # Ensure block_height_delta is between 1 and 0xFFFF (inclusive)
    dup neq.0 assert.err=ERR_TX_INVALID_EXPIRATION_DELTA
    dup push.EXPIRY_UPPER_LIMIT lt assert.err=ERR_TX_INVALID_EXPIRATION_DELTA
    # => [block_height_delta]

    exec.get_block_number add
    # => [absolute_expiration_num]

    # Load the current stored delta from memory
    dup exec.memory::get_expiration_block_num
    # => [stored_expiration_block_num, absolute_expiration_num, absolute_expiration_num]

    # Check if block_height_delta is greater
    u32lt
    if.true
        # Set new expiration delta
        exec.memory::set_expiration_block_num
    else
        drop
    end
end

#! Gets the transaction expiration delta.
#!
#! Inputs: [...]
#! Output: [block_height_delta, ...]
#!
#! Where:
#! - block_height_delta is the stored expiration time delta (1 to 0xFFFF).
export.get_expiration_delta
    exec.memory::get_expiration_block_num
    # => [stored_expiration_block_num]

    dup eq.MAX_BLOCK_NUM
    if.true
        # The delta was not set
        drop push.0
    else
        # Calculate the delta
        exec.get_block_number sub
    end
end


#! Adds a fungible asset to a note. If the note already holds an asset issued by the
#! same faucet id the two quantities are summed up and the new quantity is stored at the
#! old position in the note. In the other case, the asset is stored at the next available
#! position. Returns the pointer to the note the asset was stored at.
#! Panics if the combined quantity exceeds the maximum for fungible assets (~overflow).
#!
#! Inputs: [ASSET, note_ptr, num_of_assets, note_idx]
#! Outputs: [note_ptr]
proc.add_fungible_asset_to_note
    dup.4 exec.memory::get_output_note_asset_data_ptr
    # => [asset_ptr, ASSET, note_ptr, num_of_assets, note_idx]

    # compute the pointer at which we should stop iterating
    dup dup.7 add
    # => [end_asset_ptr, asset_ptr, ASSET, note_ptr, num_of_assets, note_idx]

    # reorganize and pad the stack, prepare for the loop
    movdn.5 movdn.5 padw dup.9
    # => [asset_ptr, 0, 0, 0, 0, ASSET, end_asset_ptr, asset_ptr, note_ptr, num_of_assets, note_idx]

    # compute the loop latch
    dup dup.10 neq
    # => [latch, asset_ptr, 0, 0, 0, 0, ASSET, end_asset_ptr, asset_ptr, note_ptr, num_of_assets, note_idx]

    while.true
        mem_loadw
        # => [STORED_ASSET, ASSET, end_asset_ptr, asset_ptr, note_ptr, num_of_assets, note_idx]

        dup.4 eq
        # => [are_equal, 0, 0, stored_amount, ASSET, end_asset_ptr, asset_ptr, note_ptr, num_of_assets, note_idx]

        if.true
            # add the asset quantity, we don't overflow here, bc both ASSETs are valid.
            movup.2 movup.6 add
            # => [updated_amount, 0, 0, faucet_id, 0, 0, end_asset_ptr, asset_ptr, note_ptr, num_of_assets, note_idx]

            # check that we don't overflow bc we use lte
            dup exec.asset::get_fungible_asset_max_amount lte assert.err=ERR_NOTE_FUNGIBLE_MAX_AMOUNT_EXCEEDED
            # => [updated_amount, 0, 0, faucet_id, 0, 0, end_asset_ptr, asset_ptr, note_ptr, num_of_assets, note_idx]

            # prepare stack to store the "updated" ASSET'' with the new quantity
            movdn.5
            # => [0, 0, ASSET'', end_asset_ptr, asset_ptr, note_ptr, num_of_assets, note_idx]

            # decrease num_of_assets by 1 to offset incrementing it later
            movup.9 sub.1 movdn.9
            # => [0, 0, ASSET'', end_asset_ptr, asset_ptr, note_ptr, num_of_assets - 1, note_idx]

            # end the loop we add 0's to the stack to have the correct number of elements
            push.0.0 dup.9 push.0
            # => [0, asset_ptr, 0, 0, 0, 0, ASSET'', end_asset_ptr, asset_ptr, note_ptr, num_of_assets - 1, note_idx]
        else
            # => [0, 0, stored_amount, ASSET, end_asset_ptr, asset_ptr, note_ptr, num_of_assets, note_idx]

            # drop ASSETs and increment the asset pointer
            movup.2 drop push.0.0 movup.9 add.1 dup movdn.10
            # => [asset_ptr + 1, 0, 0, 0, 0, ASSET, end_asset_ptr, asset_ptr + 1, note_ptr, num_of_assets, note_idx]

            # check if we reached the end of the loop
            dup dup.10 neq
        end
    end
    # => [asset_ptr, 0, 0, 0, 0, ASSET, end_asset_ptr, asset_ptr, note_ptr, num_of_assets, note_idx]
    # prepare stack for storing the ASSET
    movdn.4 dropw
    # => [asset_ptr, ASSET, end_asset_ptr, asset_ptr, note_ptr, num_of_assets, note_idx]

    # Store the fungible asset, either the combined ASSET or the new ASSET
    mem_storew dropw drop drop
    # => [note_ptr, num_of_assets, note_idx]

    # increase the number of assets in the note
    swap add.1 dup.1 exec.memory::set_output_note_num_assets
    # => [note_ptr, note_idx]
end

#! Builds the stack into the NOTE_METADATA word, encoding the note type and execution hint
#! into a single element.
#!
#! Inputs:  [tag, aux, note_type, execution_hint]
#! Outputs: [NOTE_METADATA]
proc.build_note_metadata
    # validate the note type
    # NOTE: encrypted notes are currently unsupported
    dup.2 push.PRIVATE_NOTE eq dup.3 push.PUBLIC_NOTE eq or assert.err=ERR_NOTE_INVALID_TYPE
    # => [tag, aux, note_type, execution_hint]

    # copy data to validate the tag
    dup.2 push.PUBLIC_NOTE dup.1 dup.3
    # => [tag, note_type, public_note, note_type, tag, aux, note_type, execution_hint]

    u32assert.err=ERR_NOTE_TAG_MUST_BE_U32
    # => [tag, note_type, public_note, note_type, tag, aux, note_type, execution_hint]

    # enforce the note type depending on the tag' bits
    u32shr.30 push.ALL_NOTE_TYPES_ALLOWED eq cdrop assert_eq.err=ERR_NOTE_INVALID_NOTE_TYPE_FOR_NOTE_TAG_PREFIX
    # => [tag, aux, note_type, execution_hint]

    # encode note_type and execution_hint into a single element
    movup.3 movup.3 push.TWO_POW_38 mul add movup.2
    # => [aux, encoded_type_and_ex_hint, tag]

    # add sender account ID to metadata
    exec.account::get_id
    # => [sender_acct_id, aux, encoded_type_and_ex_hint, tag]

    movdn.2
    # => [NOTE_METADATA]
end

#! Creates a new note and returns the index of the note.
#!
#! Inputs: [tag, aux, note_type, execution_hint, RECIPIENT]
#! Outputs: [note_idx]
#!
#! - tag is the tag to be included in the note.
#! - aux is the auxiliary metadata to be included in the note.
#! - note_type is the type of the note.
#! - execution_hint is the execution hint of the note.
#! - RECIPIENT defines spend conditions for the note.
#! - note_idx is the index of the crated note.
#!
#! Panics if:
#! - the note_type is not valid.
#! - the note_tag is not an u32.
#! - if note_tag starts with anything but 0b11 and note_type is not public.
#! - the number of output notes exceeds the maximum limit of 1024.
export.create_note
    push.20983 drop                                     # TODO: remove line, see miden-vm/#1122
    emit.NOTE_BEFORE_CREATED_EVENT

    exec.build_note_metadata
    # => [NOTE_METADATA, RECIPIENT]

    # get the index for the next note to be created and increment counter
    exec.increment_num_output_notes dup movdn.9
    # => [note_idx, NOTE_METADATA, RECIPIENT, note_idx]

    # get a pointer to the memory address at which the note will be stored
    exec.memory::get_output_note_ptr
    # => [note_ptr, NOTE_METADATA, RECIPIENT, note_idx]

    movdn.4
    # => [NOTE_METADATA, note_ptr, RECIPIENT, note_idx]

    # emit event to signal that a new note is created
    push.21067 drop                                     # TODO: remove line, see miden-vm/#1122
    emit.NOTE_AFTER_CREATED_EVENT

    # set the metadata for the output note
    dup.4 exec.memory::set_output_note_metadata dropw
    # => [note_ptr, RECIPIENT, note_idx]

    # set the RECIPIENT for the output note
    exec.memory::set_output_note_recipient dropw
    # => [note_idx]
end

#! Adds the ASSET to the note specified by the index.
#!
#! Inputs: [note_idx, ASSET]
#! Outputs: [note_idx]
#!
#! - note_idx is the index of the the note to which the asset is added.
#! - ASSET can be a fungible or non-fungible asset.
#!
#! Panics if:
#! - the ASSET is malformed (e.g., invalid faucet ID).
#! - the max amount of fungible assets is exceeded.
#! - the non-fungible asset already exists in the note.
#! - the total number of ASSETs exceeds the maximum of 256.
export.add_asset_to_note
    # check if the note exists, it must be within [0, num_of_notes]
    dup exec.memory::get_num_output_notes lte assert.err=ERR_NOTE_INVALID_INDEX
    # => [note_idx, ASSET]

    # get a pointer to the memory address of the note at which the asset will be stored
    dup movdn.5 exec.memory::get_output_note_ptr
    # => [note_ptr, ASSET, note_idx]

    # get current num of assets
    dup exec.memory::get_output_note_num_assets movdn.5
    # => [note_ptr, ASSET, num_of_assets, note_idx]

    # validate the ASSET
    movdn.4 exec.asset::validate_asset
    # => [ASSET, note_ptr, num_of_assets, note_idx]

    # emit event to signal that a new asset is going to be added to the note.
    push.21169 drop                                     # TODO: remove line, see miden-vm/#1122
    emit.NOTE_BEFORE_ADD_ASSET_EVENT
    # => [ASSET, note_ptr]

    # Check if ASSET to add is fungible
    exec.asset::is_fungible_asset
    # => [is_fungible_asset?, ASSET, note_ptr, num_of_assets, note_idx]

    if.true
    # ASSET to add is fungible
        exec.add_fungible_asset_to_note
        # => [note_ptr, note_idx]
    else
    # ASSET to add is non-fungible
        exec.add_non_fungible_asset_to_note
        # => [note_ptr, note_idx]
    end
    # => [note_ptr, note_idx]

    # emit event to signal that a new asset was added to the note.
    push.21277 drop                                     # TODO: remove line, see miden-vm/#1122
    emit.NOTE_AFTER_ADD_ASSET_EVENT

    drop
    # => [note_idx]
end