@layerzerolabs/lz-sui-sdk-v2/src/modules/ptb-builders/endpoint-ptb-builder.ts

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

import { ModuleManager } from '../../module-manager'
import { ObjectOptions } from '../../types'
import { asAddress, asBool, asU32, executeSimulate } from '../../utils'

export const EndpointPtbBuilderErrorCode = {
    // EndpointPtbBuilder related errors (matching endpoint_ptb_builder.move)
    EndpointPtbBuilder_EBuilderNotFound: 1,
    EndpointPtbBuilder_EBuilderRegistered: 2,
    EndpointPtbBuilder_EBuilderUnsupported: 3,
    EndpointPtbBuilder_EInvalidBounds: 4,
    EndpointPtbBuilder_EInvalidBuilderAddress: 5,
    EndpointPtbBuilder_EInvalidLibrary: 6,
    EndpointPtbBuilder_EUnauthorized: 7,
} as const

const ModuleName = 'endpoint_ptb_builder'

export class EndpointPtbBuilder {
    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 ===

    /**
     * Creates a transaction to register a message library PTB builder
     * @param tx - The transaction to add the move call to
     * @param builderInfo - PTB builder info result from getPtbBuilderInfoMoveCall or transaction argument
     */
    registerMsglibPtbBuilderMoveCall(tx: Transaction, builderInfo: TransactionArgument): void {
        tx.moveCall({
            target: this.#target('register_msglib_ptb_builder'),
            arguments: [
                tx.object(this.objects.endpointPtbBuilder),
                tx.object(this.objects.endpointPtbBuilderAdminCap),
                tx.object(this.objects.endpointV2),
                builderInfo,
            ],
        })
    }

    /**
     * Creates a transaction to set the default message library PTB builder
     * @param tx - The transaction to add the move call to
     * @param messageLib - Target message library address
     * @param ptbBuilder - PTB builder address
     */
    setDefaultMsglibPtbBuilderMoveCall(
        tx: Transaction,
        messageLib: string | TransactionArgument,
        ptbBuilder: string | TransactionArgument
    ): void {
        tx.moveCall({
            target: this.#target('set_default_msglib_ptb_builder'),
            arguments: [
                tx.object(this.objects.endpointPtbBuilder),
                tx.object(this.objects.endpointPtbBuilderAdminCap),
                asAddress(tx, messageLib),
                asAddress(tx, ptbBuilder),
            ],
        })
    }

    /**
     * Creates a transaction to set message library PTB builder for an OApp
     * @param tx - The transaction to add the move call to
     * @param call - Call object for setting msglib PTB builder
     */
    setMsglibPtbBuilderMoveCall(tx: Transaction, call: TransactionArgument): void {
        tx.moveCall({
            target: this.#target('set_msglib_ptb_builder'),
            arguments: [tx.object(this.objects.endpointPtbBuilder), call],
        })
    }

    // ======== PTB Building Functions ========

    /**
     * Build quote PTB by call for fee calculation
     * @param tx - The transaction to add the move call to
     * @param endpointQuoteCall - Endpoint quote call result
     * @returns Transaction result containing the quote PTB
     */
    buildQuotePtbByCallMoveCall(tx: Transaction, endpointQuoteCall: TransactionArgument): TransactionResult {
        return tx.moveCall({
            target: this.#target('build_quote_ptb_by_call'),
            arguments: [
                tx.object(this.objects.endpointPtbBuilder),
                tx.object(this.objects.endpointV2),
                endpointQuoteCall,
            ],
        })
    }

    /**
     * Build send PTB by call for message transmission
     * @param tx - The transaction to add the move call to
     * @param endpointSendCall - Endpoint send call result
     * @returns Transaction result containing the send PTB
     */
    buildSendPtbByCallMoveCall(tx: Transaction, endpointSendCall: TransactionArgument): TransactionResult {
        return tx.moveCall({
            target: this.#target('build_send_ptb_by_call'),
            arguments: [
                tx.object(this.objects.endpointPtbBuilder),
                tx.object(this.objects.endpointV2),
                endpointSendCall,
            ],
        })
    }

    /**
     * Build set config PTB by call for configuration updates
     * @param tx - The transaction to add the move call to
     * @param messageLibSetConfigCall - Endpoint set config call result
     * @returns Transaction result containing the set config PTB
     */
    buildSetConfigPtbByCallMoveCall(tx: Transaction, messageLibSetConfigCall: TransactionArgument): TransactionResult {
        return tx.moveCall({
            target: this.#target('build_set_config_ptb_by_call'),
            arguments: [tx.object(this.objects.endpointPtbBuilder), messageLibSetConfigCall],
        })
    }

    /**
     * Build quote PTB for on-chain composition
     * @param tx - The transaction to add the move call to
     * @param sender - Sender address
     * @param dstEid - Destination endpoint ID
     * @returns Transaction result containing the quote PTB
     */
    buildQuotePtbMoveCall(
        tx: Transaction,
        sender: string | TransactionArgument,
        dstEid: number | TransactionArgument
    ): TransactionResult {
        return tx.moveCall({
            target: this.#target('build_quote_ptb'),
            arguments: [
                tx.object(this.objects.endpointPtbBuilder),
                tx.object(this.objects.endpointV2),
                asAddress(tx, sender),
                asU32(tx, dstEid),
            ],
        })
    }

    /**
     * Build send PTB for on-chain composition
     * @param tx - The transaction to add the move call to
     * @param sender - Sender address
     * @param dstEid - Destination endpoint ID
     * @param refund - Whether to include refund step
     * @returns Transaction result containing the send PTB
     */
    buildSendPtbMoveCall(
        tx: Transaction,
        sender: string | TransactionArgument,
        dstEid: number | TransactionArgument,
        refund: boolean | TransactionArgument
    ): TransactionResult {
        return tx.moveCall({
            target: this.#target('build_send_ptb'),
            arguments: [
                tx.object(this.objects.endpointPtbBuilder),
                tx.object(this.objects.endpointV2),
                asAddress(tx, sender),
                asU32(tx, dstEid),
                asBool(tx, refund),
            ],
        })
    }

    /**
     * Build set config PTB for on-chain composition
     * @param tx - The transaction to add the move call to
     * @param sender - Sender address
     * @param lib - Message library address
     * @returns Transaction result containing the set config PTB
     */
    buildSetConfigPtbMoveCall(
        tx: Transaction,
        sender: string | TransactionArgument,
        lib: string | TransactionArgument
    ): TransactionResult {
        return tx.moveCall({
            target: this.#target('build_set_config_ptb'),
            arguments: [
                tx.object(this.objects.endpointPtbBuilder),
                tx.object(this.objects.endpointV2),
                asAddress(tx, sender),
                asAddress(tx, lib),
            ],
        })
    }

    // ======== Call ID Functions ========

    /**
     * Get endpoint quote call ID
     * @param tx - The transaction to add the move call to
     * @returns Transaction result containing the endpoint quote call ID
     */
    endpointQuoteCallIdMoveCall(tx: Transaction): TransactionResult {
        return tx.moveCall({
            target: this.#target('endpoint_quote_call_id'),
            arguments: [],
        })
    }

    /**
     * Get endpoint quote call ID as a hex string
     * @returns Promise<string> - The endpoint quote call ID in hex format
     */
    async endpointQuoteCallId(): Promise<string> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.endpointQuoteCallIdMoveCall(tx)
            },
            (result) => Buffer.from(bcs.vector(bcs.u8()).parse(result[0].value)).toString('hex')
        )
    }

    /**
     * Get endpoint send call ID
     * @param tx - The transaction to add the move call to
     * @returns Transaction result containing the endpoint send call ID
     */
    endpointSendCallIdMoveCall(tx: Transaction): TransactionResult {
        return tx.moveCall({
            target: this.#target('endpoint_send_call_id'),
            arguments: [],
        })
    }

    /**
     * Get endpoint send call ID as a hex string
     * @returns Promise<string> - The endpoint send call ID in hex format
     */
    async endpointSendCallId(): Promise<string> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.endpointSendCallIdMoveCall(tx)
            },
            (result) => Buffer.from(bcs.vector(bcs.u8()).parse(result[0].value)).toString('hex')
        )
    }

    /**
     * Get message lib quote call ID
     * @param tx - The transaction to add the move call to
     * @returns Transaction result containing the message lib quote call ID
     */
    messageLibQuoteCallIdMoveCall(tx: Transaction): TransactionResult {
        return tx.moveCall({
            target: this.#target('message_lib_quote_call_id'),
            arguments: [],
        })
    }

    /**
     * Get message lib quote call ID as a hex string
     * @returns Promise<string> - The message lib quote call ID in hex format
     */
    async messageLibQuoteCallId(): Promise<string> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.messageLibQuoteCallIdMoveCall(tx)
            },
            (result) => Buffer.from(bcs.vector(bcs.u8()).parse(result[0].value)).toString('hex')
        )
    }

    /**
     * Get message lib send call ID
     * @param tx - The transaction to add the move call to
     * @returns Transaction result containing the message lib send call ID
     */
    messageLibSendCallIdMoveCall(tx: Transaction): TransactionResult {
        return tx.moveCall({
            target: this.#target('message_lib_send_call_id'),
            arguments: [],
        })
    }

    /**
     * Get message lib send call ID as a hex string
     * @returns Promise<string> - The message lib send call ID in hex format
     */
    async messageLibSendCallId(): Promise<string> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.messageLibSendCallIdMoveCall(tx)
            },
            (result) => Buffer.from(bcs.vector(bcs.u8()).parse(result[0].value)).toString('hex')
        )
    }

    /**
     * Get message lib set config call ID
     * @param tx - The transaction to add the move call to
     * @returns Transaction result containing the message lib set config call ID
     */
    messageLibSetConfigCallIdMoveCall(tx: Transaction): TransactionResult {
        return tx.moveCall({
            target: this.#target('message_lib_set_config_call_id'),
            arguments: [],
        })
    }

    /**
     * Get message lib set config call ID as a hex string
     * @returns Promise<string> - The message lib set config call ID in hex format
     */
    async messageLibSetConfigCallId(): Promise<string> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.messageLibSetConfigCallIdMoveCall(tx)
            },
            (result) => Buffer.from(bcs.vector(bcs.u8()).parse(result[0].value)).toString('hex')
        )
    }

    // ======== View Functions ========

    /**
     * Get the default message library PTB builder for a given library
     * @param tx - The transaction to add the move call to
     * @param lib - Message library address
     * @returns Transaction result containing the default PTB builder address
     */
    getDefaultMsglibPtbBuilderMoveCall(tx: Transaction, lib: string | TransactionArgument): TransactionResult {
        return tx.moveCall({
            target: this.#target('get_default_msglib_ptb_builder'),
            arguments: [tx.object(this.objects.endpointPtbBuilder), asAddress(tx, lib)],
        })
    }

    /**
     * Get the default message library PTB builder
     * @param lib - Message library address
     * @returns Promise<string> - The default PTB builder address
     */
    async getDefaultMsglibPtbBuilder(lib: string): Promise<string> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.getDefaultMsglibPtbBuilderMoveCall(tx, lib)
            },
            (result) => bcs.Address.parse(result[0].value)
        )
    }

    /**
     * Get the OApp-specific message library PTB builder
     * @param tx - The transaction to add the move call to
     * @param oapp - OApp address or transaction argument
     * @param lib - Message library address or transaction argument
     * @returns Transaction result containing the OApp-specific PTB builder address
     */
    getOappMsglibPtbBuilderMoveCall(
        tx: Transaction,
        oapp: string | TransactionArgument,
        lib: string | TransactionArgument
    ): TransactionResult {
        return tx.moveCall({
            target: this.#target('get_oapp_msglib_ptb_builder'),
            arguments: [tx.object(this.objects.endpointPtbBuilder), asAddress(tx, oapp), asAddress(tx, lib)],
        })
    }

    /**
     * Get the OApp-specific message library PTB builder
     * @param oapp - OApp address
     * @param lib - Message library address
     * @returns Promise<string> - The OApp PTB builder address
     */
    async getOappMsglibPtbBuilder(oapp: string, lib: string): Promise<string> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.getOappMsglibPtbBuilderMoveCall(tx, oapp, lib)
            },
            (result) => bcs.Address.parse(result[0].value)
        )
    }

    /**
     * Get the effective message library PTB builder (OApp-specific or default)
     * @param tx - The transaction to add the move call to
     * @param oapp - OApp address
     * @param lib - Message library address
     * @returns Transaction result containing the effective PTB builder address
     */
    getEffectiveMsglibPtbBuilderMoveCall(
        tx: Transaction,
        oapp: string | TransactionArgument,
        lib: string | TransactionArgument
    ): TransactionResult {
        return tx.moveCall({
            target: this.#target('get_effective_msglib_ptb_builder'),
            arguments: [tx.object(this.objects.endpointPtbBuilder), asAddress(tx, oapp), asAddress(tx, lib)],
        })
    }

    /**
     * Get the effective message library PTB builder
     * @param oapp - OApp address
     * @param lib - Message library address
     * @returns Promise<string> - The effective PTB builder address
     */
    async getEffectiveMsglibPtbBuilder(oapp: string, lib: string): Promise<string> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.getEffectiveMsglibPtbBuilderMoveCall(tx, oapp, lib)
            },
            (result) => bcs.Address.parse(result[0].value)
        )
    }

    /**
     * Get detailed information about a registered message library PTB builder
     * @param tx - The transaction to add the move call to
     * @param builder - PTB builder address
     * @returns Transaction result containing the PTB builder information
     */
    getMsglibPtbBuilderInfoMoveCall(tx: Transaction, builder: string | TransactionArgument): TransactionResult {
        return tx.moveCall({
            target: this.#target('get_msglib_ptb_builder_info'),
            arguments: [tx.object(this.objects.endpointPtbBuilder), asAddress(tx, builder)],
        })
    }

    /**
     * Check if a message library PTB builder is registered
     * @param tx - The transaction to add the move call to
     * @param builder - PTB builder address
     * @returns Transaction result containing boolean registration status
     */
    isMsglibPtbBuilderRegisteredMoveCall(tx: Transaction, builder: string | TransactionArgument): TransactionResult {
        return tx.moveCall({
            target: this.#target('is_msglib_ptb_builder_registered'),
            arguments: [tx.object(this.objects.endpointPtbBuilder), asAddress(tx, builder)],
        })
    }

    /**
     * Checks if a PTB builder is registered
     * @param builder - PTB builder address
     * @returns Promise<boolean> - Whether the builder is registered
     */
    async isMsglibPtbBuilderRegistered(builder: string): Promise<boolean> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.isMsglibPtbBuilderRegisteredMoveCall(tx, builder)
            },
            (result) => bcs.Bool.parse(result[0].value)
        )
    }

    /**
     * Get the count of registered message library PTB builders
     * @param tx - The transaction to add the move call to
     * @returns Transaction result containing the count of registered builders
     */
    registeredMsglibPtbBuildersCountMoveCall(tx: Transaction): TransactionResult {
        return tx.moveCall({
            target: this.#target('registered_msglib_ptb_builders_count'),
            arguments: [tx.object(this.objects.endpointPtbBuilder)],
        })
    }

    /**
     * Get the total number of registered PTB builders
     * @returns Promise<number> - The count of registered builders
     */
    async registeredMsglibPtbBuildersCount(): Promise<number> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.registeredMsglibPtbBuildersCountMoveCall(tx)
            },
            (result) => Number(bcs.U64.parse(result[0].value))
        )
    }

    // === Private Functions ===

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