@layerzerolabs/lz-sui-sdk-v2/src/module-manager.ts

import { createHash } from 'crypto'

import { SuiClient } from '@mysten/sui/client'

import {
    BlockedMessageLib,
    BlockedMessageLibPtbBuilder,
    Call,
    Counter,
    DVN,
    DVNFeeLib,
    DvnLayerZero,
    DvnPtbBuilder,
    Endpoint,
    EndpointPtbBuilder,
    Executor,
    ExecutorFeeLib,
    ExecutorLayerZero,
    ExecutorPtbBuilder,
    LayerZeroViews,
    OApp,
    PackageWhitelistValidator,
    PriceFeed,
    PtbBuilder,
    SimpleMessageLib,
    SimpleMessageLibPtbBuilder,
    Treasury,
    Uln302,
    Uln302PtbBuilder,
    Utils,
    WorkerRegistry,
    Zro,
} from './modules'
import { ModuleOptions, ObjectOptions, PackageOptions } from './types'
import { Modules } from './types/modules'

/**
 * Module Manager - Centralized management for all SDK modules
 * Uses unified storage for both core modules and cached modules with options
 */
export class ModuleManager {
    public packages: PackageOptions
    public objects: ObjectOptions

    private storage = new Map<string, unknown>()

    constructor(packages: PackageOptions, objects: ObjectOptions) {
        this.packages = packages
        this.objects = objects
    }

    /**
     * Generate a cache key using hash for modules with options
     */
    private generateCacheKey(moduleName: string, options?: ModuleOptions): string {
        if (!options) {
            return `${moduleName}:default`
        }

        // Create hash from JSON string
        const jsonStr = JSON.stringify(options, Object.keys(options).sort())
        const hash = createHash('sha256').update(jsonStr).digest('hex').substring(0, 16)

        return `${moduleName}:${hash}`
    }

    private setCoreModules(modules: { [key: string]: unknown }): void {
        Object.entries(modules).forEach(([name, instance]) => {
            this.setModule(name, undefined, instance)
        })
    }

    /**
     * Get module instance (with options), throws error if not found
     */
    getModule<T>(moduleName: string, options?: ModuleOptions): T {
        const key = this.generateCacheKey(moduleName, options)
        const module = this.storage.get(key) as T | undefined
        if (module === undefined) {
            const optionsStr = options ? ` with options ${JSON.stringify(options)}` : ''
            throw new Error(`Module '${moduleName}'${optionsStr} not found`)
        }
        return module
    }

    /**
     * Store module instance (with options)
     */
    setModule<T>(moduleName: string, options: ModuleOptions | undefined, instance: T): void {
        const key = this.generateCacheKey(moduleName, options)
        this.storage.set(key, instance)
    }

    /**
     * Check if module (with options) exists
     */
    hasModule(moduleName: string, options?: ModuleOptions): boolean {
        const key = this.generateCacheKey(moduleName, options)
        return this.storage.has(key)
    }

    /**
     * Get module instance or create new one using factory function
     */
    getOrCreateModule<T>(moduleName: string, options: ModuleOptions | undefined, factory: () => T): T {
        // Check if module exists first
        if (this.hasModule(moduleName, options)) {
            return this.getModule<T>(moduleName, options)
        }

        // Create new instance and store it
        const instance = factory()
        this.setModule(moduleName, options, instance)
        return instance
    }

    /**
     * Get all core modules as an object
     */
    getAllCoreModules(): { [key: string]: unknown } {
        const result: { [key: string]: unknown } = {}
        const coreModuleKeys = Array.from(this.storage.keys()).filter((key) => key.endsWith(':default'))

        for (const key of coreModuleKeys) {
            const moduleName = key.replace(':default', '')
            result[moduleName] = this.storage.get(key)
        }

        return result
    }

    /**
     * Clear all modules and cache
     */
    clear(): void {
        this.storage.clear()
    }

    /**
     * Remove a module with specific options
     */
    removeModule(moduleName: string, options?: ModuleOptions): boolean {
        const key = this.generateCacheKey(moduleName, options)
        return this.storage.delete(key)
    }

    /**
     * Initialize all standard LayerZero modules
     * @param packages - Package addresses configuration
     * @param objects - Object addresses configuration
     * @param client - Sui client instance
     * @param context - Context object to be populated with module references
     */
    initializeCoreModules(packages: PackageOptions, objects: ObjectOptions, client: SuiClient): void {
        this.setCoreModules({
            [Modules.Endpoint]: new Endpoint(packages.endpointV2, client, objects, this),
            [Modules.SimpleMessageLib]: new SimpleMessageLib(packages.simpleMessageLib, client, objects, this),
            [Modules.BlockedMessageLib]: new BlockedMessageLib(packages.blockedMessageLib, client, objects, this),
            [Modules.Uln302]: new Uln302(packages.uln302, client, objects, this),
            [Modules.Utils]: new Utils(packages.utils, client),
            [Modules.Zro]: new Zro(packages.zro, client, this),
            [Modules.Call]: new Call(packages.call, client),
            [Modules.Treasury]: new Treasury(packages.treasury, client, objects, this),
            [Modules.LayerZeroViews]: new LayerZeroViews(packages.layerzeroViews, client, objects, this),
            [Modules.PtbBuilder]: new PtbBuilder(packages.ptbMoveCall, client),
        })
    }

    // === Core Module Getters ===

    getEndpoint(): Endpoint {
        return this.getModule(Modules.Endpoint)
    }

    getSimpleMessageLib(): SimpleMessageLib {
        return this.getModule(Modules.SimpleMessageLib)
    }

    getBlockedMessageLib(): BlockedMessageLib {
        return this.getModule(Modules.BlockedMessageLib)
    }

    getUln302(): Uln302 {
        return this.getModule(Modules.Uln302)
    }

    getUtils(): Utils {
        return this.getModule(Modules.Utils)
    }

    getZro(): Zro {
        return this.getModule(Modules.Zro)
    }

    getCall(): Call {
        return this.getModule(Modules.Call)
    }

    getTreasury(): Treasury {
        return this.getModule(Modules.Treasury)
    }

    getLayerZeroViews(): LayerZeroViews {
        return this.getModule(Modules.LayerZeroViews)
    }

    getPtbBuilder(): PtbBuilder {
        return this.getModule(Modules.PtbBuilder)
    }

    getOApp(client: SuiClient, callCapId: string, options?: ModuleOptions): OApp {
        return this.getOrCreateModule(Modules.Oapp + callCapId, options, () => {
            const packageId = options?.packageId ?? this.packages.oapp
            const objects = this.mergeObjectsOptions(options)
            return new OApp(packageId, callCapId, client, objects, this)
        })
    }

    // === Non-core Module Getters (created on-demand with caching) ===

    getCounter(client: SuiClient, options?: ModuleOptions): Counter {
        return this.getOrCreateModule(Modules.Counter, options, () => {
            const packageId = options?.packageId ?? this.packages.counterV2
            const objects = this.mergeObjectsOptions(options)
            return new Counter(packageId, client, objects, this)
        })
    }

    getExecutor(client: SuiClient, options?: ModuleOptions): Executor {
        return this.getOrCreateModule(Modules.Executor, options, () => {
            const packageId = options?.packageId ?? this.packages.executor
            const objects = this.mergeObjectsOptions(options)
            return new Executor(packageId, client, objects, this)
        })
    }

    getDvn(client: SuiClient, options?: ModuleOptions): DVN {
        return this.getOrCreateModule(Modules.Dvn, options, () => {
            const packageId = options?.packageId ?? this.packages.dvn
            const objects = this.mergeObjectsOptions(options)
            return new DVN(packageId, client, objects, this)
        })
    }

    getDvnFeeLib(client: SuiClient, options?: ModuleOptions): DVNFeeLib {
        return this.getOrCreateModule(Modules.DvnFeeLib, options, () => {
            const packageId = options?.packageId ?? this.packages.dvnFeeLib
            const objects = this.mergeObjectsOptions(options)
            return new DVNFeeLib(packageId, client, objects, this)
        })
    }

    getExecutorFeeLib(client: SuiClient, options?: ModuleOptions): ExecutorFeeLib {
        return this.getOrCreateModule(Modules.ExecutorFeeLib, options, () => {
            const packageId = options?.packageId ?? this.packages.executorFeeLib
            const objects = this.mergeObjectsOptions(options)
            return new ExecutorFeeLib(packageId, client, objects, this)
        })
    }

    getPriceFeed(client: SuiClient, options?: ModuleOptions): PriceFeed {
        return this.getOrCreateModule(Modules.PriceFeed, options, () => {
            const packageId = options?.packageId ?? this.packages.priceFeed
            const objects = this.mergeObjectsOptions(options)
            return new PriceFeed(packageId, client, objects, this)
        })
    }

    getDvnLayerZero(client: SuiClient, options?: ModuleOptions): DvnLayerZero {
        return this.getOrCreateModule(Modules.DvnLayerZero, options, () => {
            const packageId = options?.packageId ?? this.packages.dvnLayerzero
            return new DvnLayerZero(packageId, client, this)
        })
    }

    getExecutorLayerZero(client: SuiClient, options?: ModuleOptions): ExecutorLayerZero {
        return this.getOrCreateModule(Modules.ExecutorLayerZero, options, () => {
            const packageId = options?.packageId ?? this.packages.executorLayerzero
            return new ExecutorLayerZero(packageId, client, this)
        })
    }

    getDvnPtbBuilder(client: SuiClient, options?: ModuleOptions): DvnPtbBuilder {
        return this.getOrCreateModule(Modules.DvnPtbBuilder, options, () => {
            const packageId = options?.packageId ?? this.packages.dvnPtbBuilder
            const objects = this.mergeObjectsOptions(options)
            return new DvnPtbBuilder(packageId, client, objects, this)
        })
    }

    getExecutorPtbBuilder(client: SuiClient, options?: ModuleOptions): ExecutorPtbBuilder {
        return this.getOrCreateModule(Modules.ExecutorPtbBuilder, options, () => {
            const packageId = options?.packageId ?? this.packages.executorPtbBuilder
            const objects = this.mergeObjectsOptions(options)
            return new ExecutorPtbBuilder(packageId, client, objects, this)
        })
    }

    getPackageWhitelistValidator(client: SuiClient, options?: ModuleOptions): PackageWhitelistValidator {
        return this.getOrCreateModule(Modules.PackageWhitelistValidator, options, () => {
            const packageId = options?.packageId ?? this.packages.packageWhitelistValidator
            const objects = this.mergeObjectsOptions(options)
            return new PackageWhitelistValidator(packageId, client, objects, this)
        })
    }

    getUln302PtbBuilder(client: SuiClient, options?: ModuleOptions): Uln302PtbBuilder {
        return this.getOrCreateModule(Modules.Uln302PtbBuilder, options, () => {
            const packageId = options?.packageId ?? this.packages.uln302PtbBuilder
            const objects = this.mergeObjectsOptions(options)
            return new Uln302PtbBuilder(packageId, client, objects, this)
        })
    }

    getEndpointPtbBuilder(client: SuiClient, options?: ModuleOptions): EndpointPtbBuilder {
        return this.getOrCreateModule(Modules.EndpointPtbBuilder, options, () => {
            const packageId = options?.packageId ?? this.packages.endpointPtbBuilder
            const objects = this.mergeObjectsOptions(options)
            return new EndpointPtbBuilder(packageId, client, objects, this)
        })
    }

    getSimpleMessageLibPtbBuilder(client: SuiClient, options?: ModuleOptions): SimpleMessageLibPtbBuilder {
        return this.getOrCreateModule(Modules.SimpleMessageLibPtbBuilder, options, () => {
            const packageId = options?.packageId ?? this.packages.simpleMessageLibPtbBuilder
            const objects = this.mergeObjectsOptions(options)
            return new SimpleMessageLibPtbBuilder(packageId, client, objects, this)
        })
    }
    getBlockedMessageLibPtbBuilder(client: SuiClient, options?: ModuleOptions): BlockedMessageLibPtbBuilder {
        return this.getOrCreateModule(Modules.BlockedMessageLibPtbBuilder, options, () => {
            const packageId = options?.packageId ?? this.packages.blockedMessageLibPtbBuilder
            const objects = this.mergeObjectsOptions(options)
            return new BlockedMessageLibPtbBuilder(packageId, client, objects, this)
        })
    }

    getWorkerRegistry(client: SuiClient, options?: ModuleOptions): WorkerRegistry {
        return this.getOrCreateModule(Modules.WorkerRegistry, options, () => {
            const packageId = options?.packageId ?? this.packages.workerRegistry
            const objects = this.mergeObjectsOptions(options)
            return new WorkerRegistry(packageId, client, objects, this)
        })
    }

    // === Private Utility Methods ===

    /**
     * Merge objects configuration with options, avoiding unnecessary object spreading
     */
    private mergeObjectsOptions(options?: ModuleOptions): ObjectOptions {
        // If no options or options.objects is empty, return original objects
        if (!options?.objects || Object.keys(options.objects).length === 0) {
            return this.objects
        }

        // Only perform object spreading when there are actual properties to merge
        return { ...this.objects, ...options.objects }
    }
}