use.std::collections::smt

use.kernel::account
use.kernel::asset
use.kernel::memory

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

# The get_balance procedure can only be called on a fungible faucet
const.ERR_VAULT_GET_BALANCE_PROC_CAN_ONLY_BE_CALLED_ON_FUNGIBLE_FAUCET=0x00020018

# The has_non_fungible_asset procedure can only be called on a non-fungible faucet
const.ERR_VAULT_HAS_NON_FUNGIBLE_ASSET_PROC_CAN_BE_CALLED_ONLY_WITH_NON_FUNGIBLE_ASSET=0x00020019

# Adding the fungible asset to the vault would exceed the max amount of 9223372036854775807
const.ERR_VAULT_FUNGIBLE_MAX_AMOUNT_EXCEEDED=0x0002001A

# Failed to add fungible asset to the asset vault due to the initial value being invalid
const.ERR_VAULT_ADD_FUNGIBLE_ASSET_FAILED_INITIAL_VALUE_INVALID=0x0002001B

# The non-fungible asset already exists in the asset vault
const.ERR_VAULT_NON_FUNGIBLE_ASSET_ALREADY_EXISTS=0x0002001C

# Failed to remove the fungible asset from the vault since the amount of the asset in the vault is less than the amount to remove
const.ERR_VAULT_FUNGIBLE_ASSET_AMOUNT_LESS_THAN_AMOUNT_TO_WITHDRAW=0x0002001D

# Failed to remove fungible asset from the asset vault due to the initial value being invalid
const.ERR_VAULT_REMOVE_FUNGIBLE_ASSET_FAILED_INITIAL_VALUE_INVALID=0x0002001E

# Failed to remove non-existent non-fungible asset from the vault
const.ERR_VAULT_NON_FUNGIBLE_ASSET_TO_REMOVE_NOT_FOUND=0x0002001F

# ACCESSORS
# =================================================================================================

#! Returns the balance of a fungible asset associated with a faucet_id.
#! Panics if the asset is not a fungible asset.
#!
#! Stack: [faucet_id, vault_root_ptr]
#! Output: [balance]
#!
#! - vault_root_ptr is a pointer to the memory location at which the vault root is stored.
#! - faucet_id is the faucet id of the fungible asset of interest.
#! - balance is the vault balance of the fungible asset.
export.get_balance
    # assert that the faucet id is a fungible faucet
    dup exec.account::is_fungible_faucet assert.err=ERR_VAULT_GET_BALANCE_PROC_CAN_ONLY_BE_CALLED_ON_FUNGIBLE_FAUCET
    # => [faucet_id, vault_root_ptr]

    # get the asset vault root
    padw movup.5 mem_loadw
    # => [ASSET_VAULT_ROOT, faucet_id]

    # prepare the key for fungible asset lookup (pad least significant elements with zeros)
    push.0 push.0 push.0 movup.7
    # => [faucet_id, 0, 0, 0, ASSET_VAULT_ROOT]

    # lookup asset
    exec.smt::get swapw dropw
    # => [ASSET]

    # extract asset balance (ASSET[0])
    drop drop drop
    # => [balance]
end

#! Returns a boolean indicating whether the non-fungible asset is present in the vault.
#! Panics if the ASSET is a fungible asset.
#!
#! Stack: [ASSET, vault_root_ptr]
#! Output: [has_asset]
#!
#! - vault_root_ptr is a pointer to the memory location at which the vault root is stored.
#! - ASSET is the non-fungible asset of interest
#! - has_asset is a boolean indicating whether the account vault has the asset of interest
export.has_non_fungible_asset
    # check if the asset is a non-fungible asset
    exec.asset::is_non_fungible_asset assert.err=ERR_VAULT_HAS_NON_FUNGIBLE_ASSET_PROC_CAN_BE_CALLED_ONLY_WITH_NON_FUNGIBLE_ASSET
    # => [ASSET, vault_root_ptr]

    # prepare the stack to read non-fungible asset from vault
    padw movup.8 mem_loadw swapw
    # => [ASSET, ACCT_VAULT_ROOT]

    # lookup asset
    exec.smt::get swapw dropw
    # => [ASSET]

    # compare with EMPTY_WORD to asses if the asset exists in the vault
    padw eqw not
    # => [has_asset, PAD, ASSET]

    # organize the stack for return
    movdn.4 dropw movdn.4 dropw
    # => [has_asset]
end

# ADD ASSET
# =================================================================================================

#! Add the specified fungible asset to the vault.  If the vault already contains an asset
#! issued by the same faucet, the amounts are added together.
#!
#! Panics:
#! - If the total value of assets is greater than or equal to 2^63.
#!
#! Stack: [ASSET, vault_root_ptr]
#! Output: [ASSET']
#!
#! - vault_root_ptr is a pointer to the memory location at which the vault root is stored.
#! - ASSET is the fungible asset to add to the vault.
#! - ASSET' is the total fungible asset in the account vault after ASSET was added to it.
export.add_fungible_asset
    push.0 movdn.3 dup movdn.4
    # => [ASSET_KEY, faucet_id, amount, vault_root_ptr]

    # get the asset vault root and read the vault asset value using the `push_smtpeek` decorator.
    # To account for the edge case in which CUR_VAULT_VALUE is an EMPTY_WORD, we replace the most
    # significant element with the faucet_id to construct the CUR_ASSET.
    padw dup.10 mem_loadw swapw
    adv.push_smtpeek push.15329 drop            # TODO: remove line, see miden-vm/#1122
    adv_loadw swapw dupw.1 drop movup.11
    # => [CUR_ASSET, VAULT_ROOT, CUR_VAULT_VALUE, amount, vault_root_ptr]

    # arrange elements
    movup.3 movup.12 dup
    # => [amount, amount, cur_amount, faucet_id, 0, 0, VAULT_ROOT, CUR_VAULT_VALUE, vault_root_ptr]

    # compute max_amount - cur_amount
    exec.asset::get_fungible_asset_max_amount dup.3 sub
    # => [(max_amount - cur_amount), amount, amount, cur_amount, faucet_id, 0, 0, VAULT_ROOT, CUR_VAULT_VALUE, vault_root_ptr]

    # assert amount + cur_amount < max_amount
    lte assert.err=ERR_VAULT_FUNGIBLE_MAX_AMOUNT_EXCEEDED
    # => [amount, cur_amount, faucet_id, 0, 0, VAULT_ROOT, CUR_VAULT_VALUE, vault_root_ptr]

    # add asset amounts
    add movdn.3
    # => [ASSET', VAULT_ROOT, CUR_VAULT_VALUE, vault_root_ptr]

    # prepare the stack to insert the asset into the vault
    dupw movdnw.3 dupw movup.3 drop push.0 movdn.3 swapw
    # => [ASSET', KEY, VAULT_ROOT, CUR_VAULT_VALUE, ASSET', vault_root_ptr]

    # update asset in vault and assert the old value is equivalent to the value provided via the
    # decorator
    exec.smt::set movupw.2 assert_eqw.err=ERR_VAULT_ADD_FUNGIBLE_ASSET_FAILED_INITIAL_VALUE_INVALID
    # => [VAULT_ROOT', ASSET', vault_root_ptr]

    # update the vault root
    movup.8 mem_storew dropw
    # => [ASSET']
end

#! Add the specified non-fungible asset to the vault.
#!
#! Panics:
#! - If the vault already contains the same non-fungible asset.
#!
#! Stack: [ASSET, vault_root_ptr]
#! Output: [ASSET]
#!
#! - vault_root_ptr is a pointer to the memory location at which the vault root is stored.
#! - ASSET is the non-fungible asset that is added to the vault.
export.add_non_fungible_asset
    # prepare the stack to insert the asset into the vault
    dup.4 movdn.5 dupw padw movup.12 mem_loadw swapw dupw
    # => [ASSET, ASSET, VAULT_ROOT, ASSET, vault_root_ptr]

    # insert asset into vault
    exec.smt::set
    # => [OLD_VAL, VAULT_ROOT', ASSET, vault_root_ptr]

    # Assert old value was empty
    padw assert_eqw.err=ERR_VAULT_NON_FUNGIBLE_ASSET_ALREADY_EXISTS
    # => [VAULT_ROOT', ASSET, vault_root_ptr]

    # update the vault root
    movup.8 mem_storew dropw
    # => [ASSET]
end

#! Add the specified asset to the vault.
#!
#! Panics:
#! - If the asset is not valid.
#! - If the total value of two fungible assets is greater than or equal to 2^63.
#! - If the vault already contains the same non-fungible asset.
#!
#! Stack: [ASSET, vault_root_ptr]
#! Output: [ASSET']
#!
#! - ASSET is the asset that is added to the vault.
#! - vault_root_ptr is a pointer to the memory location at which the vault root is stored.
#! - ASSET' final asset in the account vault defined as follows:
#!   - If ASSET is a non-fungible asset, then ASSET' is the same as ASSET.
#!   - If ASSET is a fungible asset, then ASSET' is the total fungible asset in the account vault
#!     after ASSET was added to it.
export.add_asset
    # check if the asset is a fungible asset
    exec.asset::is_fungible_asset
    # => [is_fungible_asset, ASSET]

    # add the asset to the asset vault
    if.true
        # validate the fungible asset
        exec.asset::validate_fungible_asset
        # => [ASSET]

        exec.add_fungible_asset
        # => [ASSET']
    else
        # validate the non-fungible asset
        exec.asset::validate_non_fungible_asset
        # => [ASSET]

        exec.add_non_fungible_asset
        # => [ASSET']
    end
end

# REMOVE ASSET
# =================================================================================================

#! Remove the specified fungible asset from the vault.
#!
#! Panics:
#! - The amount of the asset in the vault is less than the amount to be removed.
#!
#! Stack: [ASSET, vault_root_ptr]
#! Output: [ASSET]
#!
#! - ASSET is the fungible asset to remove from the vault.
#! - vault_root_ptr is a pointer to the memory location at which the vault root is stored.
export.remove_fungible_asset
    dupw push.0 movdn.3 dup movdn.4
    # => [ASSET_KEY, faucet_id, amount, ASSET, vault_root_ptr]

    # get the asset vault root and read the vault asset value using the `push_smtpeek` decorator
    # To account for the edge case in which CUR_VAULT_VALUE is an EMPTY_WORD, we replace the most
    # significant element with the faucet_id to construct the CUR_ASSET.
    padw dup.14 mem_loadw swapw
    adv.push_smtpeek push.15413 drop            # TODO: remove line, see miden-vm/#1122
    adv_loadw dupw movdnw.2 drop movup.11
    # => [CUR_ASSET, VAULT_ROOT, CUR_VAULT_VALUE, amount, ASSET, vault_root_ptr]

    # arrange elements
    movup.3 movup.12 dup dup.2
    # => [cur_amount, amount, amount, cur_amount, faucet_id, 0, 0, VAULT_ROOT, CUR_VAULT_VALUE, ASSET, vault_root_ptr]

    # assert amount <= cur_amount
    lte assert.err=ERR_VAULT_FUNGIBLE_ASSET_AMOUNT_LESS_THAN_AMOUNT_TO_WITHDRAW
    # => [amount, cur_amount, faucet_id, 0, 0, VAULT_ROOT, CUR_VAULT_VALUE, ASSET, vault_root_ptr]

    # asset amount + cur_amount < max_amount
    sub
    # => [new_amount, faucet_id, 0, 0, VAULT_ROOT, CUR_VAULT_VALUE, ASSET, vault_root_ptr]

    # => check if the asset amount is zero
    dup eq.0
    # => [is_zero, new_amount, faucet_id, 0, 0, VAULT_ROOT, CUR_VAULT_VALUE, ASSET, vault_root_ptr]

    if.true
        # fungible asset empty - insert EMPTY_WORD in vault
        movdn.3 padw
        # => [EMPTY_WORD, ASSET_KEY, VAULT_ROOT, CUR_VAULT_VALUE, ASSET, vault_root_ptr]
    else
        # fungible asset not empty - update asset in vault
        movdn.3 dupw movup.3 drop push.0 movdn.3 swapw
        # => [NEW_ASSET, ASSET_KEY, VAULT_ROOT, CUR_VAULT_VALUE, ASSET, vault_root_ptr]
    end

    # update asset in vault and assert the old value is equivalent to the value provided via the
    # decorator
    exec.smt::set movupw.2 assert_eqw.err=ERR_VAULT_REMOVE_FUNGIBLE_ASSET_FAILED_INITIAL_VALUE_INVALID
    # => [VAULT_ROOT', ASSET, vault_root_ptr]

    # update the vault root
    movup.8 mem_storew dropw
    # => [ASSET]
end

#! Remove the specified non-fungible asset from the vault.
#!
#! Panics:
#! - The non-fungible asset is not found in the vault.
#!
#! Stack: [ASSET, vault_root_ptr]
#! Output: [ASSET]
#!
#! - ASSET is the non-fungible asset to remove from the vault.
#! - vault_root_ptr is a pointer to the memory location at which the vault root is stored.
export.remove_non_fungible_asset
    # prepare the stack to insert an EMPTY_WORD into the vault at key associated with the
    # non-fungible asset
    dup.4 movdn.5 dupw padw movup.12 mem_loadw swapw padw
    # => [EMPTY_WORD, ASSET, VAULT_ROOT, ASSET, vault_root_ptr]

    # update asset in vault
    exec.smt::set
    # => [OLD_VAL, VAULT_ROOT', ASSET, vault_root_ptr]

    # Assert old value was not empty (we only need to check ASSET[1] which is the faucet id)
    drop drop eq.0 assertz.err=ERR_VAULT_NON_FUNGIBLE_ASSET_TO_REMOVE_NOT_FOUND drop
    # => [VAULT_ROOT', ASSET, vault_root_ptr]

    # update the vault root
    movup.8 mem_storew dropw
    # => [ASSET]
end

#! Remove the specified asset from the vault.
#!
#! Panics:
#! - The fungible asset is not found in the vault.
#! - The amount of the fungible asset in the vault is less than the amount to be removed.
#! - The non-fungible asset is not found in the vault.
#!
#! Stack: [ASSET, vault_root_ptr]
#! Output: [ASSET]
#!
#! - ASSET is the asset to remove from the vault.
#! - vault_root_ptr is a pointer to the memory location at which the vault root is stored.
export.remove_asset
    # check if the asset is a fungible asset
    exec.asset::is_fungible_asset
    # => [is_fungible_asset, ASSET, vault_root_ptr]

    # remove the asset from the asset vault
    if.true
        exec.remove_fungible_asset
        # => [ASSET]
    else
        exec.remove_non_fungible_asset
        # => [ASSET]
    end
end