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

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

import { OAppInfoV1Bcs, VectorMoveCallBCS } from '../../bcs'
import {
    Argument,
    BuilderPlaceholderInfo,
    DEFAULT_SIMULATION_TIMES,
    LzComposeVersion,
    MoveCall,
    OAppInfoVersion,
    SimulateResult,
} from '../../types'
import { simulateTransaction } from '../../utils/transaction'
import { normalizeSuiPackageId } from '../../utils/type-name'

const MOVE_CALL_MODULE_NAME = 'move_call'

export const PtbBuilderErrorCode = {
    // MoveCallsBuilder related errors (matching move_calls_builder.move)
    MoveCallsBuilder_EInvalidMoveCallResult: 1,
    MoveCallsBuilder_EResultIDNotFound: 2,

    // Argument related errors (matching argument.move)
    Argument_EInvalidArgument: 1,
} as const

export class PtbBuilder {
    public packageId: string
    public readonly client: SuiClient

    constructor(packageId: string, client: SuiClient) {
        this.packageId = packageId
        this.client = client
    }

    /**
     * Simulate a transaction and decode the resulting move calls
     * @param tx - The transaction to simulate
     * @param sender - Optional sender address for simulation context
     * @returns Promise resolving to array of decoded move calls
     */
    async simulatePtb(tx: Transaction, sender?: string): Promise<MoveCall[]> {
        const ptbCallsResult = await simulateTransaction(this.client, tx, sender)
        return this.#decodeMoveCalls(ptbCallsResult[0])
    }

    /**
     * Simulate a LayerZero receive transaction and decode the move calls
     * Handles versioned receive data and decodes based on version
     * @param tx - The transaction to simulate
     * @param sender - Optional sender address for simulation context
     * @returns Promise resolving to array of decoded move calls
     * @throws Error if unsupported version is encountered
     */
    async simulateLzReceivePtb(tx: Transaction, sender?: string): Promise<MoveCall[]> {
        const ptbCallsResult = await simulateTransaction(this.client, tx, sender)
        const buffer = Buffer.from(bcs.vector(bcs.u8()).parse(ptbCallsResult[0].value))

        const version = buffer.readInt16BE()
        if (version === OAppInfoVersion.VERSION_1) {
            return this.#decodeMoveCallsFromOAppInfoV1(new Uint8Array(buffer.subarray(2)))
        }
        throw new Error(`Unknown version: ${version}`)
    }

    /**
     * Simulate a LayerZero compose transaction and decode the move calls
     * Handles versioned compose data and decodes based on version
     * @param tx - The transaction to simulate
     * @param sender - Optional sender address for simulation context
     * @returns Promise resolving to array of decoded move calls
     * @throws Error if unsupported version is encountered
     */
    async simulateLzComposePtb(tx: Transaction, sender?: string): Promise<MoveCall[]> {
        const ptbCallsResult = await simulateTransaction(this.client, tx, sender)
        const buffer = Buffer.from(bcs.vector(bcs.u8()).parse(ptbCallsResult[0].value))
        const version = buffer.readInt16BE()
        if (version === LzComposeVersion.VERSION_1) {
            return this.#decodeMoveCallsBytes(new Uint8Array(buffer.subarray(2)))
        }
        throw new Error(`Unknown version: ${version}`)
    }

    /**
     * Builds PTB with move-calls simulated from the transaction
     *
     * This method processes an array of move calls, handling both regular calls and builder calls
     * (which require simulation to expand into actual move calls). It ensures all object arguments
     * are properly validated for PTB compatibility.
     *
     * @param tx - The transaction to append move calls to
     * @param moveCalls - Array of move calls to process and build
     * @param resolutionIDs - Cache mapping call IDs to their transaction results for argument resolution (defaults to empty Map)
     * @param sender - Optional sender address for simulation context (defaults to undefined)
     * @param maxSimulationTimes - Maximum number of simulations allowed for builder calls (defaults to DEFAULT_SIMULATION_TIMES)
     * @param nestedResult - Array storing results from previous calls for NestedResult argument resolution (internal use, defaults to empty array)
     * @param baseOffset - Base offset for calculating nested result indices (internal use, defaults to 0)
     *
     * @returns Promise<[number, MoveCall[]]> - [moveCallCount, finalMoveCalls] tuple
     *
     * @throws Error if simulation limit is exceeded, nested results are unavailable, or objects are not PTB-compatible
     */
    async buildPtb(
        tx: Transaction,
        moveCalls: MoveCall[],
        resolutionIDs = new Map<string, TransactionResult>(), // ID -> TransactionResult
        sender: string | undefined = undefined,
        maxSimulationTimes = DEFAULT_SIMULATION_TIMES,
        // -- below are internal use only --
        nestedResult: TransactionResult[] = [],
        baseOffset = 0
    ): Promise<[number, MoveCall[]]> {
        const finalMoveCalls: MoveCall[] = [] // This array collects all move calls for validation
        const moveCallCount = await this.#buildMoveCalls(
            tx,
            moveCalls,
            resolutionIDs,
            sender,
            maxSimulationTimes,
            finalMoveCalls,
            nestedResult,
            baseOffset
        )

        return [moveCallCount, finalMoveCalls]
    }

    /**
     * Internal method to recursively build and process move calls
     * Handles both regular and builder calls with simulation and argument resolution
     * @param tx - The transaction to add move calls to
     * @param moveCalls - Array of move calls to process
     * @param resolutionIDs - Map for resolving call ID arguments
     * @param sender - Optional sender address for simulation
     * @param remainingSimulation - Remaining simulation attempts allowed
     * @param finalMoveCalls - Array collecting all final move calls
     * @param nestedResult - Array of transaction results for nested argument resolution
     * @param baseOffset - Base offset for calculating nested result indices
     * @returns Promise resolving to the number of move calls processed
     * @private
     */
    async #buildMoveCalls(
        tx: Transaction,
        moveCalls: MoveCall[],
        resolutionIDs: Map<string, TransactionResult>, // ID -> TransactionResult
        sender: string | undefined = undefined,
        // -- below are internal use only --
        remainingSimulation: number,
        finalMoveCalls: MoveCall[],
        nestedResult: TransactionResult[] = [],
        baseOffset = 0
    ): Promise<number> {
        if (!moveCalls.length) return 0
        let builderMoveCallCount = 0 // current builder move_calls count
        const placeholderInfos: BuilderPlaceholderInfo[] = []
        for (let currentMoveCallIndex = 0; currentMoveCallIndex < moveCalls.length; currentMoveCallIndex++) {
            const moveCall = moveCalls[currentMoveCallIndex]
            // Create a copy to avoid mutating the original
            const processedMoveCall = {
                ...moveCall,
                arguments: moveCall.arguments.map((arg) => {
                    // Adjust the index by the base offset and the placeholder count
                    if ('NestedResult' in arg) {
                        const [callIndex, resultIndex] = arg.NestedResult
                        const newCallIndex = this.#calculateOffset(callIndex, baseOffset, placeholderInfos)
                        const result: Argument = { NestedResult: [newCallIndex, resultIndex] }
                        return result
                    }
                    return arg
                }),
            }

            if (!moveCall.is_builder_call) {
                finalMoveCalls.push(processedMoveCall)
                this.#appendMoveCall(tx, processedMoveCall, resolutionIDs, nestedResult, true)
            } else {
                builderMoveCallCount++

                if (remainingSimulation <= 0) {
                    throw new Error('remainingSimulation is not enough')
                }
                remainingSimulation--
                // The simluate tx is simluated to get the actual move calls of the contract to be replaced in the next layer
                const simulateTx = Transaction.from(tx)
                this.#appendMoveCall(simulateTx, processedMoveCall, resolutionIDs, nestedResult, false)
                const newMoveCalls = await this.simulatePtb(simulateTx, sender)

                // Replace the placeholder move calls with the actual move calls
                const placeholderMoveCallCount = await this.#buildMoveCalls(
                    tx,
                    newMoveCalls,
                    resolutionIDs,
                    sender,
                    remainingSimulation,
                    finalMoveCalls,
                    nestedResult,
                    this.#calculateOffset(currentMoveCallIndex, baseOffset, placeholderInfos)
                )
                placeholderInfos.push({ index: currentMoveCallIndex, moveCallCount: placeholderMoveCallCount })
            }
        }
        const placeholderMoveCallCount = placeholderInfos.reduce((acc, item) => acc + item.moveCallCount, 0)
        return moveCalls.length - builderMoveCallCount + placeholderMoveCallCount
    }

    /**
     * Append a move call to the transaction with argument resolution
     * Handles different argument types and records results for future reference
     * @param tx - The transaction to add the move call to
     * @param moveCall - The move call to append
     * @param resolutionIDs - Map for resolving call ID arguments
     * @param nestedResult - Array to store transaction results
     * @param directCall - Whether this is a direct call (affects result recording)
     * @private
     */
    #appendMoveCall(
        tx: Transaction,
        moveCall: MoveCall,
        resolutionIDs: Map<string, TransactionResult>, // ID -> TransactionResult
        nestedResult: TransactionResult[],
        directCall: boolean
    ): void {
        const moveCallParam = {
            target: `${moveCall.function.package}::${moveCall.function.module_name}::${moveCall.function.name}`,
            arguments: moveCall.arguments.map((arg) => {
                if ('Object' in arg) {
                    return tx.object(arg.Object)
                } else if ('Pure' in arg) {
                    return tx.pure(arg.Pure)
                } else if ('NestedResult' in arg) {
                    const [callIndex, resultIndex] = arg.NestedResult
                    const moveCallResult = nestedResult[callIndex]
                    // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition,@typescript-eslint/strict-boolean-expressions
                    if (!moveCallResult || typeof moveCallResult[resultIndex] === 'undefined') {
                        throw new Error(`NestedResult resultIndex ${resultIndex} not available in call ${callIndex}`)
                    }
                    return moveCallResult[resultIndex]
                } else if ('ID' in arg) {
                    // Replace the call parameters with the result from the call cache
                    const result = resolutionIDs.get(Buffer.from(arg.ID).toString('hex'))
                    if (result === undefined) {
                        throw new Error(
                            `Call substitution not found for "${Buffer.from(arg.ID).toString('hex')}" in cache`
                        )
                    }
                    return result
                }
                throw new Error(`Unknown argument variant: ${JSON.stringify(arg)}`)
            }),
            typeArguments: moveCall.type_arguments,
        }

        if (directCall) {
            const result = tx.moveCall(moveCallParam)
            // The nestedResult index is actually equal to the move_call's nested result index in the move_call_builder
            nestedResult.push(result)
            // Record output calls for downstream parameter replacement if there are output calls
            if (moveCall.result_ids.length > 0) {
                for (let i = 0; i < moveCall.result_ids.length; i++) {
                    const outputID = moveCall.result_ids[i]
                    // Validate that the result has enough outputs
                    if (i >= result.length) {
                        throw new Error(
                            `Move call result doesn't have output at index ${i} for call "${Buffer.from(outputID).toString('hex')}"`
                        )
                    }
                    const resultID = Buffer.from(outputID).toString('hex')
                    if (resolutionIDs.has(resultID)) {
                        throw new Error(`Move call result already exists for call "${resultID}"`)
                    }
                    resolutionIDs.set(resultID, result[i] as unknown as TransactionResult)
                }
            }
        } else {
            tx.moveCall(moveCallParam)
        }
    }

    /**
     * Decode move calls from a simulation result
     * @param viewResult - The simulation result containing move calls
     * @returns Array of decoded move calls
     * @throws Error if the result type doesn't match expected format
     * @private
     */
    #decodeMoveCalls(viewResult: SimulateResult): MoveCall[] {
        const { value, type } = viewResult
        const vectorCallType =
            `vector<${normalizeSuiPackageId(this.packageId, true, true)}::${MOVE_CALL_MODULE_NAME}::MoveCall>`.toLowerCase()
        if (type.toLowerCase() !== vectorCallType) {
            throw new Error(`not match the type: ${type.toLowerCase()} - expected ${vectorCallType}`)
        }
        return this.#decodeMoveCallsBytes(value)
    }

    #decodeMoveCallsFromOAppInfoV1(bytes: Uint8Array): MoveCall[] {
        const oappInfo = OAppInfoV1Bcs.parse(bytes)
        return this.#decodeMoveCallsBytes(new Uint8Array(oappInfo.lz_receive_info).subarray(2)) // remove the version prefix
    }

    /**
     * Decode move calls from raw bytes
     * @param bytes - The raw bytes containing encoded move calls
     * @returns Array of decoded move calls
     * @private
     */
    #decodeMoveCallsBytes(bytes: Uint8Array): MoveCall[] {
        const moveCalls = VectorMoveCallBCS.parse(bytes)
        return moveCalls.map((moveCall) => ({
            function: {
                package: moveCall.function.package,
                module_name: moveCall.function.module_name,
                name: moveCall.function.name,
            },
            arguments: moveCall.arguments.map((arg) => {
                if (arg.ID !== undefined) {
                    return { ID: arg.ID }
                } else if (arg.Object !== undefined) {
                    return { Object: arg.Object }
                } else if (arg.Pure !== undefined) {
                    return { Pure: new Uint8Array(arg.Pure) }
                } else {
                    // Must be NestedResult since we've handled other cases
                    return {
                        NestedResult: [arg.NestedResult.call_index, arg.NestedResult.result_index] as [number, number],
                    }
                }
            }),
            type_arguments: moveCall.type_arguments,
            result_ids: moveCall.result_ids.map((id) => Buffer.from(id)),
            is_builder_call: moveCall.is_builder_call,
        })) as MoveCall[]
    }

    /**
     * Calculate the offset for nested result indices accounting for builder expansions
     * @param index - The original index to calculate offset for
     * @param baseOffset - The base offset to add
     * @param placeholderInfos - Information about placeholder expansions
     * @returns The calculated offset index
     * @private
     */
    #calculateOffset(index: number, baseOffset: number, placeholderInfos: BuilderPlaceholderInfo[]): number {
        let placeholderCount = 0
        let placeholderMoveCallCount = 0
        for (const placeholder of placeholderInfos) {
            // builder move_call will not return any nested result, so only check index > placeholder.index
            if (index > placeholder.index) {
                placeholderMoveCallCount += placeholder.moveCallCount
                placeholderCount++
            }
        }
        return index - placeholderCount + placeholderMoveCallCount + baseOffset
    }
}