@layerzerolabs/lz-sui-sdk-v2/src/modules/workers/treasury.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, asU64, executeSimulate } from '../../utils'

const MODULE_NAME = 'treasury'

export const TreasuryErrorCode = {
    // Treasury related errors (matching treasury.move)
    TREASURY_EInvalidFeeRecipient: 1,
    TREASURY_EInvalidNativeFeeBp: 2,
    TREASURY_EZroNotEnabled: 3,
} as const

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

    /**
     * Set fee recipient address for treasury (admin only)
     * @param tx - The transaction to add the move call to
     * @param feeRecipient - The new fee recipient address or transaction argument
     */
    setFeeRecipientMoveCall(tx: Transaction, feeRecipient: string | TransactionArgument): void {
        tx.moveCall({
            target: this.#target('set_fee_recipient'),
            arguments: [
                tx.object(this.objects.treasury),
                tx.object(this.objects.treasuryAdminCap),
                asAddress(tx, feeRecipient),
            ],
        })
    }

    /**
     * Set native fee basis points for treasury (admin only)
     * @param tx - The transaction to add the move call to
     * @param nativeFeeBp - The native fee in basis points or transaction argument
     */
    setNativeFeeBpMoveCall(tx: Transaction, nativeFeeBp: bigint | number | string | TransactionArgument): void {
        tx.moveCall({
            target: this.#target('set_native_fee_bp'),
            arguments: [
                tx.object(this.objects.treasury),
                tx.object(this.objects.treasuryAdminCap),
                asU64(tx, nativeFeeBp),
            ],
        })
    }

    /**
     * Set ZRO token fee for treasury (admin only)
     * @param tx - The transaction to add the move call to
     * @param zroFee - The ZRO token fee amount or transaction argument
     */
    setZroFeeMoveCall(tx: Transaction, zroFee: bigint | number | string | TransactionArgument): void {
        tx.moveCall({
            target: this.#target('set_zro_fee'),
            arguments: [tx.object(this.objects.treasury), tx.object(this.objects.treasuryAdminCap), asU64(tx, zroFee)],
        })
    }

    /**
     * Enable or disable ZRO token fees (admin only)
     * @param tx - The transaction to add the move call to
     * @param zroEnabled - Whether ZRO token fees are enabled or transaction argument
     */
    setZroEnabledMoveCall(tx: Transaction, zroEnabled: boolean | TransactionArgument): void {
        tx.moveCall({
            target: this.#target('set_zro_enabled'),
            arguments: [
                tx.object(this.objects.treasury),
                tx.object(this.objects.treasuryAdminCap),
                asBool(tx, zroEnabled),
            ],
        })
    }

    /**
     * Enable or disable treasury fees (admin only)
     * @param tx - The transaction to add the move call to
     * @param feeEnabled - Whether treasury fees are enabled or transaction argument
     */
    setFeeEnabledMoveCall(tx: Transaction, feeEnabled: boolean | TransactionArgument): void {
        tx.moveCall({
            target: this.#target('set_fee_enabled'),
            arguments: [
                tx.object(this.objects.treasury),
                tx.object(this.objects.treasuryAdminCap),
                asBool(tx, feeEnabled),
            ],
        })
    }

    // === View Functions ===

    /**
     * Calculate treasury fees for a transaction
     * @param tx - The transaction to add the move call to
     * @param totalNativeFee - Total native fee amount or transaction argument
     * @param payInZro - Whether to pay in ZRO tokens or transaction argument
     * @returns Transaction result containing calculated fees
     */
    getFeeMoveCall(
        tx: Transaction,
        totalNativeFee: bigint | number | string | TransactionArgument,
        payInZro: boolean | TransactionArgument
    ): TransactionResult {
        return tx.moveCall({
            target: this.#target('get_fee'),
            arguments: [tx.object(this.objects.treasury), asU64(tx, totalNativeFee), asBool(tx, payInZro)],
        })
    }

    /**
     * Calculate treasury fees for a transaction
     * @param totalNativeFee - Total native fee amount
     * @param payInZro - Whether to pay in ZRO tokens
     * @returns Promise<[bigint, bigint]> - Tuple of [nativeFee, zroFee]
     */
    async getFee(totalNativeFee: bigint, payInZro: boolean): Promise<[bigint, bigint]> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.getFeeMoveCall(tx, totalNativeFee, payInZro)
            },
            (result) => {
                // Parse the tuple (u64, u64) returned by get_fee
                const nativeFee = BigInt(bcs.U64.parse(result[0].value))
                const zroFee = BigInt(bcs.U64.parse(result[1].value))
                return [nativeFee, zroFee]
            }
        )
    }

    /**
     * Get treasury fee recipient address
     * @param tx - The transaction to add the move call to
     * @returns Transaction result containing the fee recipient address
     */
    feeRecipientMoveCall(tx: Transaction): TransactionResult {
        return tx.moveCall({
            target: this.#target('fee_recipient'),
            arguments: [tx.object(this.objects.treasury)],
        })
    }

    /**
     * Get treasury fee recipient address
     * @returns Promise<string> - The fee recipient address
     */
    async feeRecipient(): Promise<string> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.feeRecipientMoveCall(tx)
            },
            (result) => bcs.Address.parse(result[0].value)
        )
    }

    /**
     * Get native fee basis points
     * @param tx - The transaction to add the move call to
     * @returns Transaction result containing the native fee basis points
     */
    nativeFeeBpMoveCall(tx: Transaction): TransactionResult {
        return tx.moveCall({
            target: this.#target('native_fee_bp'),
            arguments: [tx.object(this.objects.treasury)],
        })
    }

    /**
     * Get native fee basis points
     * @returns Promise<bigint> - The native fee in basis points
     */
    async nativeFeeBp(): Promise<bigint> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.nativeFeeBpMoveCall(tx)
            },
            (result) => BigInt(bcs.U64.parse(result[0].value))
        )
    }

    /**
     * Get ZRO token fee amount
     * @param tx - The transaction to add the move call to
     * @returns Transaction result containing the ZRO fee
     */
    zroFeeMoveCall(tx: Transaction): TransactionResult {
        return tx.moveCall({
            target: this.#target('zro_fee'),
            arguments: [tx.object(this.objects.treasury)],
        })
    }

    /**
     * Get ZRO token fee amount
     * @returns Promise<bigint> - The ZRO token fee
     */
    async zroFee(): Promise<bigint> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.zroFeeMoveCall(tx)
            },
            (result) => BigInt(bcs.U64.parse(result[0].value))
        )
    }

    /**
     * Check if ZRO token fees are enabled
     * @param tx - The transaction to add the move call to
     * @returns Transaction result containing the ZRO enabled status
     */
    zroEnabledMoveCall(tx: Transaction): TransactionResult {
        return tx.moveCall({
            target: this.#target('zro_enabled'),
            arguments: [tx.object(this.objects.treasury)],
        })
    }

    /**
     * Check if ZRO token fees are enabled
     * @returns Promise<boolean> - True if ZRO fees are enabled
     */
    async zroEnabled(): Promise<boolean> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.zroEnabledMoveCall(tx)
            },
            (result) => bcs.Bool.parse(result[0].value)
        )
    }

    /**
     * Check if treasury fees are enabled
     * @param tx - The transaction to add the move call to
     * @returns Transaction result containing the fee enabled status
     */
    feeEnabledMoveCall(tx: Transaction): TransactionResult {
        return tx.moveCall({
            target: this.#target('fee_enabled'),
            arguments: [tx.object(this.objects.treasury)],
        })
    }

    /**
     * Check if treasury fees are enabled
     * @returns Promise<boolean> - True if treasury fees are enabled
     */
    async feeEnabled(): Promise<boolean> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.feeEnabledMoveCall(tx)
            },
            (result) => bcs.Bool.parse(result[0].value)
        )
    }

    // === 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 MODULE_NAME)
     * @returns The full module path for the move call
     * @private
     */
    #target(name: string, module_name: string = MODULE_NAME): string {
        return `${this.packageId}::${module_name}::${name}`
    }
}