@layerzerolabs/lz-sui-sdk-v2/src/modules/workers/worker-registry.ts

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, executeSimulate } from '../../utils'

const MODULE_NAME = 'worker_registry'

export const WorkerRegistryErrorCode = {
    // WorkerRegistry related errors (matching worker_registry.move)
    EWorkerInfoInvalid: 1,
    EWorkerInfoNotFound: 2,
} as const

/**
 * WorkerRegistry provides access to the worker registry contract functionality
 * This module maintains mappings between worker addresses and their associated information.
 */
export class WorkerRegistry {
    public packageId: string
    public readonly client: SuiClient
    private readonly objects: ObjectOptions

    /**
     * Create a new WorkerRegistry instance
     * @param packageId - The package ID of the worker registry contract
     * @param client - The Sui client instance
     * @param objects - Object options containing contract object IDs
     * @param moduleManager - Module manager for handling dependencies
     */
    constructor(
        packageId: string,
        client: SuiClient,
        objects: ObjectOptions,
        private readonly moduleManager: ModuleManager
    ) {
        this.packageId = packageId
        this.client = client
        this.objects = objects
    }

    // === View Functions ===

    /**
     * Get worker info for a specific worker address (as a move call)
     * This creates a move call that can be used in a PTB (Programmable Transaction Block)
     * @param tx - The transaction to add the move call to
     * @param worker - The worker address or transaction argument
     * @returns TransactionResult - The move call result containing worker info bytes
     * @throws Will throw an error if worker info is not found (EWorkerInfoNotFound)
     */
    getWorkerInfoMoveCall(tx: Transaction, worker: string | TransactionArgument): TransactionResult {
        return tx.moveCall({
            target: this.#target('get_worker_info'),
            arguments: [tx.object(this.objects.workerRegistry), asAddress(tx, worker)],
        })
    }

    /**
     * Get worker info for a specific worker address
     * Executes the get_worker_info function and returns the worker information bytes
     * @param worker - The worker address to get info for
     * @returns Promise<Uint8Array> - The worker information bytes containing encoded worker data
     * @throws Will throw an error if worker info is not found (EWorkerInfoNotFound)
     */
    async getWorkerInfo(worker: string): Promise<Uint8Array> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.getWorkerInfoMoveCall(tx, worker)
            },
            (result) => {
                return result[0].value
            }
        )
    }

    /**
     * Check if worker info exists for a specific worker address
     * This is a convenience method that catches EWorkerInfoNotFound errors
     * @param worker - The worker address to check
     * @returns Promise<boolean> - True if worker info exists, false if not found
     * @throws Will re-throw any errors other than EWorkerInfoNotFound
     */
    async hasWorkerInfo(worker: string): Promise<boolean> {
        try {
            await this.getWorkerInfo(worker)
            return true
        } catch (error) {
            // If the error is EWorkerInfoNotFound, return false
            if (error instanceof Error && error.message.includes('EWorkerInfoNotFound')) {
                return false
            }
            // Re-throw other errors
            throw error
        }
    }

    // === Private Functions ===

    /**
     * Generate the target string for move calls
     * @param functionName - The function name to call
     * @returns The fully qualified target string for the move call
     */
    #target(functionName: string): string {
        return `${this.packageId}::${MODULE_NAME}::${functionName}`
    }
}