@layerzerolabs/lz-sui-sdk-v2/src/modules/oapps/counter.ts

import { bcs } from '@mysten/sui/bcs'
import { SuiClient } from '@mysten/sui/client'
import { Transaction, TransactionArgument, TransactionResult } from '@mysten/sui/transactions'

import { EndpointId } from '@layerzerolabs/lz-definitions'

import { ModuleManager } from '../../module-manager'
import { MessagingFee, ObjectOptions } from '../../types'
import {
    IPTBValidator,
    asAddress,
    asBool,
    asBytes,
    asBytes32,
    asU32,
    asU64,
    asU8,
    executeSimulate,
    isTransactionArgument,
} from '../../utils'

const MODULE_NAME = 'counter'

export const CounterErrorCode = {
    // Counter related errors (matching counter.move)
    Counter_EInvalidMsgType: 0,
    Counter_EInvalidValue: 1,
    Counter_EInvalidNonce: 2,
    Counter_EOnlyEndpoint: 3,
    Counter_ESelfComposeOnly: 4,
    Counter_EInvalidOApp: 5,

    // OptionsBuilder related errors (matching options_builder.move)
    OptionsBuilder_EInvalidSize: 1,
} as const

export class Counter {
    public packageId: string
    public readonly client: SuiClient
    private readonly objects: ObjectOptions

    constructor(
        packageId: string,
        client: SuiClient,
        objects: ObjectOptions,
        private readonly moduleManager: ModuleManager
    ) {
        this.packageId = packageId
        this.client = client
        this.objects = objects
    }

    // === Set Functions ===

    /**
     * Initialize counter with LayerZero receive and compose information
     * @param tx - The transaction to add the move call to
     * @param lzReceiveInfo - LayerZero receive information transaction argument
     * @param lzComposeInfo - LayerZero compose information transaction argument
     */
    initCounterMoveCall(tx: Transaction, lzReceiveInfo: TransactionArgument, lzComposeInfo: TransactionArgument): void {
        tx.moveCall({
            target: this.#target('init_counter'),
            arguments: [
                tx.object(this.objects.counter),
                tx.object(this.objects.counterOapp),
                tx.object(this.objects.counterAdminCap),
                tx.object(this.objects.endpointV2),
                lzReceiveInfo,
                lzComposeInfo,
            ],
        })
    }

    /**
     * Quote messaging fees for sending counter increment
     * @param dstEid - Destination endpoint ID
     * @param msgType - Message type (SEND or SEND_AND_CALL)
     * @param options - Execution options as bytes
     * @param payInZero - Whether to pay in ZRO tokens
     * @returns Promise<MessagingFee> - The calculated messaging fees
     */
    async quote(
        dstEid: EndpointId | TransactionArgument,
        msgType: number | TransactionArgument,
        options: Uint8Array | TransactionArgument,
        payInZero: boolean | TransactionArgument,
        validators?: IPTBValidator[]
    ): Promise<MessagingFee> {
        const tx = new Transaction()
        const quoteCall = tx.moveCall({
            target: this.#target('quote'),
            arguments: [
                tx.object(this.objects.counter),
                tx.object(this.objects.counterOapp),
                asU32(tx, dstEid),
                asU8(tx, msgType),
                asBytes(tx, options),
                asBool(tx, payInZero),
            ],
        })
        return this.moduleManager.getEndpoint().quote(tx, quoteCall, undefined, validators)
    }

    /**
     * Increment counter on destination chain
     * @param tx - The transaction to add the move call to
     * @param sender - Sender address for ZRO token operations
     * @param dstEid - Destination endpoint ID
     * @param msgType - Message type (SEND or SEND_AND_CALL)
     * @param options - Execution options as bytes
     * @param nativeFee - Native token fee amount
     * @param zroFee - ZRO token fee amount
     * @param refundAddress - Address for fee refunds
     */
    async incrementMoveCall(
        tx: Transaction,
        sender: string,
        dstEid: EndpointId | TransactionArgument,
        msgType: number | TransactionArgument,
        options: Uint8Array | TransactionArgument,
        nativeFee: bigint | TransactionArgument,
        zroFee: bigint | TransactionArgument,
        refundAddress: string | TransactionArgument,
        validators?: IPTBValidator[]
    ): Promise<void> {
        const [nativeToken] = tx.splitCoins(tx.gas, [asU64(tx, nativeFee)])
        const zroToken = isTransactionArgument(zroFee)
            ? zroFee
            : await this.moduleManager.getZro().splitOptionZroTokenMoveCall(tx, sender, zroFee)
        const incrementCall = tx.moveCall({
            target: this.#target('increment'),
            arguments: [
                tx.object(this.objects.counter),
                tx.object(this.objects.counterOapp),
                asU32(tx, dstEid),
                asU8(tx, msgType),
                asBytes(tx, options),
                nativeToken,
                zroToken,
                asAddress(tx, refundAddress),
            ],
        })
        await this.moduleManager.getEndpoint().populateSendTransaction(tx, incrementCall, sender, validators)
    }

    /**
     * Set composer information for counter
     * @param tx - The transaction to add the move call to
     * @param composerInfo - Composer information including lz_compose execution information as bytes
     */
    setComposerInfoMoveCall(tx: Transaction, composerInfo: Uint8Array | TransactionArgument): void {
        tx.moveCall({
            target: this.#target('set_composer_info'),
            arguments: [
                tx.object(this.objects.counter),
                tx.object(this.objects.counterOapp),
                tx.object(this.objects.counterAdminCap),
                tx.object(this.objects.endpointV2),
                asBytes(tx, composerInfo),
            ],
        })
    }

    // === View Functions ===

    /**
     * Get counter EID
     * @param tx - The transaction to add the move call to
     * @returns Transaction result containing the counter EID
     */
    getEidMoveCall(tx: Transaction): TransactionResult {
        return tx.moveCall({
            target: this.#target('eid'),
            arguments: [tx.object(this.objects.counter)],
        })
    }

    /**
     * Get counter EID
     * @returns Promise<number> - The counter EID
     */
    async getEid(): Promise<number> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.getEidMoveCall(tx)
            },
            (result) => bcs.U32.parse(result[0].value)
        )
    }

    /**
     * Get call capability address for counter
     * @param tx - The transaction to add the move call to
     * @returns Transaction result containing the call capability address
     */
    getCallCapAddressMoveCall(tx: Transaction): TransactionResult {
        return tx.moveCall({
            target: this.#target('call_cap_address'),
            arguments: [tx.object(this.objects.counter)],
        })
    }

    /**
     * Get call capability address for counter
     * @returns Promise<string> - The call capability address
     */
    async getCallCapAddress(): Promise<string> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.getCallCapAddressMoveCall(tx)
            },
            (result) => bcs.Address.parse(result[0].value)
        )
    }

    /**
     * Get composer address for counter
     * @param tx - The transaction to add the move call to
     * @returns Transaction result containing the composer address
     */
    getComposerAddressMoveCall(tx: Transaction): TransactionResult {
        return tx.moveCall({
            target: this.#target('composer_address'),
            arguments: [tx.object(this.objects.counter)],
        })
    }

    /**
     * Get composer address for counter
     * @returns Promise<string> - The composer address
     */
    async getComposerAddress(): Promise<string> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.getComposerAddressMoveCall(tx)
            },
            (result) => bcs.Address.parse(result[0].value)
        )
    }

    /**
     * Get current counter value
     * @param tx - The transaction to add the move call to
     * @returns Transaction result containing the counter value
     */
    getCountMoveCall(tx: Transaction): TransactionResult {
        return tx.moveCall({
            target: this.#target('get_count'),
            arguments: [tx.object(this.objects.counter)],
        })
    }

    /**
     * Get current counter value
     * @returns Promise<number> - The current counter value
     */
    async getCount(): Promise<number> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.getCountMoveCall(tx)
            },
            (result) => Number(bcs.U64.parse(result[0].value))
        )
    }

    /**
     * Get composed counter value
     * @param tx - The transaction to add the move call to
     * @returns Transaction result containing the composed counter value
     */
    getComposedCountMoveCall(tx: Transaction): TransactionResult {
        return tx.moveCall({
            target: this.#target('get_composed_count'),
            arguments: [tx.object(this.objects.counter)],
        })
    }

    /**
     * Get composed counter value
     * @returns Promise<number> - The composed counter value
     */
    async getComposedCount(): Promise<number> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.getComposedCountMoveCall(tx)
            },
            (result) => Number(bcs.U64.parse(result[0].value))
        )
    }

    /**
     * Get outbound counter value for a destination
     * @param tx - The transaction to add the move call to
     * @param dstEid - Destination endpoint ID
     * @returns Transaction result containing the outbound counter value
     */
    getOutboundCountMoveCall(tx: Transaction, dstEid: number | TransactionArgument): TransactionResult {
        return tx.moveCall({
            target: this.#target('get_outbound_count'),
            arguments: [tx.object(this.objects.counter), asU32(tx, dstEid)],
        })
    }

    /**
     * Get outbound counter value for a destination
     * @param dstEid - Destination endpoint ID
     * @returns Promise<number> - The outbound counter value
     */
    async getOutboundCount(dstEid: number): Promise<number> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.getOutboundCountMoveCall(tx, dstEid)
            },
            (result) => Number(bcs.U64.parse(result[0].value))
        )
    }

    /**
     * Get inbound counter value from a source
     * @param tx - The transaction to add the move call to
     * @param srcEid - Source endpoint ID
     * @returns Transaction result containing the inbound counter value
     */
    getInboundCountMoveCall(tx: Transaction, srcEid: number | TransactionArgument): TransactionResult {
        return tx.moveCall({
            target: this.#target('get_inbound_count'),
            arguments: [tx.object(this.objects.counter), asU32(tx, srcEid)],
        })
    }

    /**
     * Get inbound counter value from a source
     * @param srcEid - Source endpoint ID
     * @returns Promise<number> - The inbound counter value
     */
    async getInboundCount(srcEid: number): Promise<number> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.getInboundCountMoveCall(tx, srcEid)
            },
            (result) => Number(bcs.U64.parse(result[0].value))
        )
    }

    nextNonceMoveCall(
        tx: Transaction,
        srcEid: number | TransactionArgument,
        sender: Uint8Array | TransactionArgument
    ): TransactionResult {
        return tx.moveCall({
            target: this.#target('next_nonce'),
            arguments: [
                tx.object(this.objects.counter),
                asU32(tx, srcEid),
                asBytes32(tx, sender, this.moduleManager.getUtils()),
            ],
        })
    }

    async nextNonce(srcEid: number, sender: Uint8Array): Promise<number> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.nextNonceMoveCall(tx, srcEid, sender)
            },
            (result) => Number(bcs.U64.parse(result[0].value))
        )
    }

    getMaxReceivedNonceMoveCall(
        tx: Transaction,
        srcEid: number | TransactionArgument,
        sender: Uint8Array | TransactionArgument
    ): TransactionResult {
        return tx.moveCall({
            target: this.#target('get_max_received_nonce'),
            arguments: [
                tx.object(this.objects.counter),
                asU32(tx, srcEid),
                asBytes32(tx, sender, this.moduleManager.getUtils()),
            ],
        })
    }

    async getMaxReceivedNonce(srcEid: number, sender: Uint8Array): Promise<number> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.getMaxReceivedNonceMoveCall(tx, srcEid, sender)
            },
            (result) => Number(bcs.U64.parse(result[0].value))
        )
    }

    isOrderedNonceMoveCall(tx: Transaction): TransactionResult {
        return tx.moveCall({
            target: this.#target('is_ordered_nonce'),
            arguments: [tx.object(this.objects.counter)],
        })
    }

    async isOrderedNonce(): Promise<boolean> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.isOrderedNonceMoveCall(tx)
            },
            (result) => bcs.Bool.parse(result[0].value)
        )
    }

    // === PTB Builder Functions ===

    lzReceiveInfoMoveCall(tx: Transaction): TransactionResult {
        return tx.moveCall({
            target: this.#target('lz_receive_info', 'counter_ptb_builder'),
            arguments: [
                tx.object(this.objects.counter),
                tx.object(this.objects.counterOapp),
                tx.object(this.objects.endpointV2),
                tx.object(this.objects.endpointPtbBuilder),
            ],
        })
    }

    lzComposeInfoMoveCall(tx: Transaction): TransactionResult {
        return tx.moveCall({
            target: this.#target('lz_compose_info', 'counter_ptb_builder'),
            arguments: [
                tx.object(this.objects.counter),
                tx.object(this.objects.counterOapp),
                tx.object(this.objects.endpointV2),
                tx.object(this.objects.endpointPtbBuilder),
            ],
        })
    }

    // === Private Functions ===

    /**
     * Generate the full target path for move calls
     * @param name - The function name to call
     * @param module_name - The module name (defaults to 'counter')
     * @returns The full module path for the move call
     * @private
     */
    #target(name: string, module_name = MODULE_NAME): string {
        return `${this.packageId}::${module_name}::${name}`
    }
}