@layerzerolabs/lz-sui-sdk-v2/src/modules/workers/price-feed.ts

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

import { ArbitrumPriceExtBcs, ModelTypeBcs, PriceBcs } from '../../bcs'
import { ModuleManager } from '../../module-manager'
import { ArbitrumPriceExt, ModelType, ObjectOptions, Price } from '../../types'
import { asAddress, asBool, asObject, asU128, asU32, asU64, executeSimulate } from '../../utils'

const MODULE_NAME = 'price_feed'

export const PriceFeedErrorCode = {
    // PriceFeed related errors
    PRICE_FEED_EInvalidDenominator: 1,
    PRICE_FEED_ENoPrice: 2,
    PRICE_FEED_ENotAnOPStack: 3,
    PRICE_FEED_EOnlyPriceUpdater: 4,
    PRICE_FEED_EPriceUpdaterCapNotFound: 5,
} as const

export class PriceFeed {
    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
    }

    // === Helper Functions ===

    /**
     * Create price configuration object
     * @param tx - The transaction to add the move call to
     * @param priceRatio - Price ratio value or transaction argument
     * @param gasPriceInUnit - Gas price in unit or transaction argument
     * @param gasPerByte - Gas per byte value or transaction argument
     * @returns Transaction result containing the price object
     */
    createPriceMoveCall(
        tx: Transaction,
        priceRatio: bigint | number | string | TransactionArgument,
        gasPriceInUnit: bigint | number | string | TransactionArgument,
        gasPerByte: number | TransactionArgument
    ): TransactionResult {
        return tx.moveCall({
            target: this.#target('create_price'),
            arguments: [asU128(tx, priceRatio), asU64(tx, gasPriceInUnit), asU32(tx, gasPerByte)],
        })
    }

    /**
     * Create Arbitrum price extension object
     * @param tx - The transaction to add the move call to
     * @param gasPerL2Tx - Gas per L2 transaction or transaction argument
     * @param gasPerL1CallDataByte - Gas per L1 call data byte or transaction argument
     * @returns Transaction result containing the Arbitrum price extension object
     */
    createArbitrumPriceExtMoveCall(
        tx: Transaction,
        gasPerL2Tx: bigint | number | string | TransactionArgument,
        gasPerL1CallDataByte: number | TransactionArgument
    ): TransactionResult {
        return tx.moveCall({
            target: this.#target('create_arbitrum_price_ext'),
            arguments: [asU64(tx, gasPerL2Tx), asU32(tx, gasPerL1CallDataByte)],
        })
    }

    /**
     * Get price model type move call result
     * @param tx - The transaction to add the move call to
     * @param modelType - The price model type enum value
     * @returns Transaction result containing the model type
     */
    getModelTypeMoveCall(tx: Transaction, modelType: ModelType): TransactionResult {
        switch (modelType) {
            case ModelType.DEFAULT:
                return tx.moveCall({
                    target: this.#target('model_type_default'),
                    arguments: [],
                })
            case ModelType.ARB_STACK:
                return tx.moveCall({
                    target: this.#target('model_type_arbitrum'),
                    arguments: [],
                })
            case ModelType.OP_STACK:
                return tx.moveCall({
                    target: this.#target('model_type_optimism'),
                    arguments: [],
                })
            default:
                throw new Error(`Invalid model type: ${JSON.stringify(modelType)}`)
        }
    }

    // === Set Functions ===

    /**
     * Set price updater role for an address (admin only)
     * Note: This function will automatically create a price updater capability for new updaters
     * @param tx - The transaction to add the move call to
     * @param updater - The updater address or transaction argument
     * @param active - Whether to activate or deactivate the updater role or transaction argument
     */
    setPriceUpdaterMoveCall(
        tx: Transaction,
        updater: string | TransactionArgument,
        active: boolean | TransactionArgument
    ): void {
        tx.moveCall({
            target: this.#target('set_price_updater'),
            arguments: [
                tx.object(this.objects.priceFeed),
                tx.object(this.objects.priceFeedOwnerCap),
                asAddress(tx, updater),
                asBool(tx, active),
            ],
        })
    }

    /**
     * Set price ratio denominator for price calculations (admin only)
     * Note: denominator must be greater than 0, otherwise the transaction will fail
     * @param tx - The transaction to add the move call to
     * @param denominator - The price ratio denominator value or transaction argument (must be > 0)
     */
    setPriceRatioDenominatorMoveCall(
        tx: Transaction,
        denominator: bigint | number | string | TransactionArgument
    ): void {
        tx.moveCall({
            target: this.#target('set_price_ratio_denominator'),
            arguments: [
                tx.object(this.objects.priceFeed),
                tx.object(this.objects.priceFeedOwnerCap),
                asU128(tx, denominator),
            ],
        })
    }

    /**
     * Set Arbitrum compression percentage (admin only)
     * @param tx - The transaction to add the move call to
     * @param compressionPercent - The compression percentage for Arbitrum or transaction argument
     */
    setArbitrumCompressionPercentMoveCall(
        tx: Transaction,
        compressionPercent: bigint | number | string | TransactionArgument
    ): void {
        tx.moveCall({
            target: this.#target('set_arbitrum_compression_percent'),
            arguments: [
                tx.object(this.objects.priceFeed),
                tx.object(this.objects.priceFeedOwnerCap),
                asU128(tx, compressionPercent),
            ],
        })
    }

    /**
     * Set price model type for a destination EID (admin only)
     * @param tx - The transaction to add the move call to
     * @param dstEid - Destination endpoint ID or transaction argument
     * @param modelType - The price model type to set
     */
    setEidToModelTypeMoveCall(tx: Transaction, dstEid: number | TransactionArgument, modelType: ModelType): void {
        const modelTypeCall = this.getModelTypeMoveCall(tx, modelType)

        tx.moveCall({
            target: this.#target('set_eid_to_model_type'),
            arguments: [
                tx.object(this.objects.priceFeed),
                tx.object(this.objects.priceFeedOwnerCap),
                asU32(tx, dstEid),
                modelTypeCall,
            ],
        })
    }

    /**
     * Set price for a destination EID (price updater capability required)
     * @param tx - The transaction to add the move call to
     * @param updaterCap - The price updater capability object or transaction argument
     * @param dstEid - Destination endpoint ID or transaction argument
     * @param price - The price configuration to set
     */
    setPriceMoveCall(
        tx: Transaction,
        updaterCap: string | TransactionArgument,
        dstEid: number | TransactionArgument,
        price: Price
    ): void {
        const priceCall = this.createPriceMoveCall(tx, price.priceRatio, price.gasPriceInUnit, price.gasPerByte)

        tx.moveCall({
            target: this.#target('set_price'),
            arguments: [tx.object(this.objects.priceFeed), asObject(tx, updaterCap), asU32(tx, dstEid), priceCall],
        })
    }

    /**
     * Set price for Arbitrum with additional extension parameters (price updater capability required)
     * @param tx - The transaction to add the move call to
     * @param updaterCap - The price updater capability object or transaction argument
     * @param dstEid - Destination endpoint ID
     * @param price - The base price configuration
     * @param arbitrumPriceExt - Additional Arbitrum-specific price parameters
     */
    setPriceForArbitrumMoveCall(
        tx: Transaction,
        updaterCap: string | TransactionArgument,
        dstEid: number | TransactionArgument,
        price: Price,
        arbitrumPriceExt: ArbitrumPriceExt
    ): void {
        const priceCall = this.createPriceMoveCall(tx, price.priceRatio, price.gasPriceInUnit, price.gasPerByte)
        const arbitrumPriceExtCall = this.createArbitrumPriceExtMoveCall(
            tx,
            arbitrumPriceExt.gasPerL2Tx,
            arbitrumPriceExt.gasPerL1CallDataByte
        )

        tx.moveCall({
            target: this.#target('set_price_for_arbitrum'),
            arguments: [
                tx.object(this.objects.priceFeed),
                asObject(tx, updaterCap),
                asU32(tx, dstEid),
                priceCall,
                arbitrumPriceExtCall,
            ],
        })
    }

    /**
     * Set native token price in USD (price updater capability required)
     * @param tx - The transaction to add the move call to
     * @param updaterCap - The price updater capability object or transaction argument
     * @param nativeTokenPriceUsd - The native token price in USD
     */
    setNativeTokenPriceUsdMoveCall(
        tx: Transaction,
        updaterCap: string | TransactionArgument,
        nativeTokenPriceUsd: bigint | number | string | TransactionArgument
    ): void {
        tx.moveCall({
            target: this.#target('set_native_token_price_usd'),
            arguments: [tx.object(this.objects.priceFeed), asObject(tx, updaterCap), asU128(tx, nativeTokenPriceUsd)],
        })
    }

    // === Witness Functions ===

    /**
     * Create a LayerZero witness for PriceFeed package whitelist registration
     * @param tx - The transaction to add the move call to
     * @returns Transaction result containing the LayerZero witness
     */
    createLayerZeroWitnessMoveCall(tx: Transaction): TransactionResult {
        return tx.moveCall({
            target: `${this.packageId}::price_feed_witness::new`,
            arguments: [],
        })
    }

    // === Worker Function ===

    /**
     * Estimate fee by endpoint ID using a call result
     * @param tx - The transaction to add the move call to
     * @param call - The call transaction result containing fee parameters
     */
    estimateFeeByEidMoveCall(tx: Transaction, call: TransactionResult): void {
        tx.moveCall({
            target: this.#target('estimate_fee_by_eid'),
            arguments: [tx.object(this.objects.priceFeed), call],
        })
    }

    // === View Functions ===

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

    /**
     * Get price updater capability address for a specific updater
     * @param tx - The transaction to add the move call to
     * @param updater - The updater address to get capability for
     * @returns Transaction result containing the price updater capability address
     */
    getPriceUpdaterCapMoveCall(tx: Transaction, updater: string | TransactionArgument): TransactionResult {
        return tx.moveCall({
            target: this.#target('get_price_updater_cap'),
            arguments: [tx.object(this.objects.priceFeed), asAddress(tx, updater)],
        })
    }

    /**
     * Get the owner capability address of this PriceFeed
     * @returns Promise<string> - The owner capability address
     */
    async ownerCap(): Promise<string> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.getOwnerCapMoveCall(tx)
            },
            (result) => bcs.Address.parse(result[0].value)
        )
    }

    /**
     * Get price updater capability address for a specific updater
     * @param updater - The updater address to get capability for
     * @returns Promise<string> - The price updater capability address
     */
    async priceUpdaterCap(updater: string): Promise<string> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.getPriceUpdaterCapMoveCall(tx, updater)
            },
            (result) => bcs.Address.parse(result[0].value)
        )
    }

    /**
     * Check if an address is a price updater
     * @param tx - The transaction to add the move call to
     * @param updater - The updater address to check
     * @returns Transaction result containing the price updater status
     */
    isPriceUpdaterMoveCall(tx: Transaction, updater: string | TransactionArgument): TransactionResult {
        return tx.moveCall({
            target: this.#target('is_price_updater'),
            arguments: [tx.object(this.objects.priceFeed), asAddress(tx, updater)],
        })
    }

    /**
     * Check if an address is a price updater
     * @param updater - The updater address to check
     * @returns Promise<boolean> - True if the address is a price updater
     */
    async isPriceUpdater(updater: string): Promise<boolean> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.isPriceUpdaterMoveCall(tx, updater)
            },
            (result) => bcs.Bool.parse(result[0].value)
        )
    }

    /**
     * Get price ratio denominator
     * @param tx - The transaction to add the move call to
     * @returns Transaction result containing the price ratio denominator
     */
    priceRatioDenominatorMoveCall(tx: Transaction): TransactionResult {
        return tx.moveCall({
            target: this.#target('get_price_ratio_denominator'),
            arguments: [tx.object(this.objects.priceFeed)],
        })
    }

    /**
     * Get price ratio denominator
     * @returns Promise<bigint> - The price ratio denominator value
     */
    async priceRatioDenominator(): Promise<bigint> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.priceRatioDenominatorMoveCall(tx)
            },
            (result) => BigInt(bcs.U128.parse(result[0].value))
        )
    }

    /**
     * Get Arbitrum compression percentage
     * @param tx - The transaction to add the move call to
     * @returns Transaction result containing the compression percentage
     */
    arbitrumCompressionPercentMoveCall(tx: Transaction): TransactionResult {
        return tx.moveCall({
            target: this.#target('get_arbitrum_compression_percent'),
            arguments: [tx.object(this.objects.priceFeed)],
        })
    }

    /**
     * Get Arbitrum compression percentage
     * @returns Promise<bigint> - The compression percentage value
     */
    async arbitrumCompressionPercent(): Promise<bigint> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.arbitrumCompressionPercentMoveCall(tx)
            },
            (result) => BigInt(bcs.U128.parse(result[0].value))
        )
    }

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

    /**
     * Get model type for a destination EID
     * @param dstEid - Destination endpoint ID
     * @returns Promise<PriceModelType> - The price model type
     */
    async modelType(dstEid: number): Promise<ModelType> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.modelTypeMoveCall(tx, dstEid)
            },
            (result) => {
                const enumValue = ModelTypeBcs.parse(result[0].value) as { $kind: string }
                switch (enumValue.$kind) {
                    case 'DEFAULT':
                        return ModelType.DEFAULT
                    case 'ARB_STACK':
                        return ModelType.ARB_STACK
                    case 'OP_STACK':
                        return ModelType.OP_STACK
                    default:
                        throw new Error(`Invalid model type: ${JSON.stringify(enumValue)}`)
                }
            }
        )
    }

    /**
     * Get native token price in USD
     * @param tx - The transaction to add the move call to
     * @returns Transaction result containing the native token price in USD
     */
    nativeTokenPriceUsdMoveCall(tx: Transaction): TransactionResult {
        return tx.moveCall({
            target: this.#target('native_token_price_usd'),
            arguments: [tx.object(this.objects.priceFeed)],
        })
    }

    /**
     * Get native token price in USD
     * @returns Promise<bigint> - The native token price in USD
     */
    async nativeTokenPriceUsd(): Promise<bigint> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.nativeTokenPriceUsdMoveCall(tx)
            },
            (result) => BigInt(bcs.U128.parse(result[0].value))
        )
    }

    /**
     * Get Arbitrum price extension parameters
     * @param tx - The transaction to add the move call to
     * @returns Transaction result containing Arbitrum price extension
     */
    arbitrumPriceExtMoveCall(tx: Transaction): TransactionResult {
        return tx.moveCall({
            target: this.#target('arbitrum_price_ext'),
            arguments: [tx.object(this.objects.priceFeed)],
        })
    }

    /**
     * Get Arbitrum price extension parameters
     * @returns Promise<ArbitrumPriceExt> - The Arbitrum price extension configuration
     */
    async arbitrumPriceExt(): Promise<ArbitrumPriceExt> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.arbitrumPriceExtMoveCall(tx)
            },
            (result) => {
                const parsed = ArbitrumPriceExtBcs.parse(result[0].value) as {
                    gas_per_l2_tx: string | number | bigint
                    gas_per_l1_call_data_byte: number
                }
                return {
                    gasPerL2Tx: BigInt(parsed.gas_per_l2_tx),
                    gasPerL1CallDataByte: parsed.gas_per_l1_call_data_byte,
                }
            }
        )
    }

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

    /**
     * Get price for a specific destination EID
     * @param dstEid - Destination endpoint ID
     * @returns Promise<Price> - The price configuration for the destination
     */
    async price(dstEid: number): Promise<Price> {
        return executeSimulate(
            this.client,
            (tx) => {
                this.priceMoveCall(tx, dstEid)
            },
            (result) => {
                const parsed = PriceBcs.parse(result[0].value) as {
                    price_ratio: string | number | bigint
                    gas_price_in_unit: string | number | bigint
                    gas_per_byte: number
                }
                return {
                    priceRatio: BigInt(parsed.price_ratio),
                    gasPriceInUnit: BigInt(parsed.gas_price_in_unit),
                    gasPerByte: parsed.gas_per_byte,
                }
            }
        )
    }

    /**
     * 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 = MODULE_NAME): string {
        return `${this.packageId}::${module_name}::${name}`
    }
}