neptune-web-wallet/src/composables/useNeptuneWallet.ts

import { useNeptuneStore } from '@/stores/neptuneStore'
import * as API from '@/api/neptuneApi'
import type {
    GenerateSeedResult,
    PayloadBuildTransaction,
    ViewKeyResult,
    WalletState,
} from '@/interface'
import initWasm, { generate_seed, address_from_seed, validate_seed_phrase } from '@neptune/wasm'

let wasmInitialized = false
let initPromise: Promise<void> | null = null

export function useNeptuneWallet() {
    const store = useNeptuneStore()

    // ===== WASM METHODS =====
    const ensureWasmInitialized = async (): Promise<void> => {
        if (wasmInitialized) {
            return
        }

        if (initPromise) {
            return initPromise
        }

        initPromise = (async () => {
            try {
                await initWasm()

                wasmInitialized = true
            } catch (err) {
                wasmInitialized = false
                const errorMsg = 'Failed to initialize Neptune WASM'
                console.error('WASM init error:', err)
                throw new Error(errorMsg)
            }
        })()

        return initPromise
    }

    const generateWallet = async (): Promise<GenerateSeedResult> => {
        try {
            await ensureWasmInitialized()

            const resultJson = generate_seed()
            const result: GenerateSeedResult = JSON.parse(resultJson)

            store.setSeedPhrase(result.seed_phrase)
            store.setReceiverId(result.receiver_identifier)

            const viewKeyResult = await getViewKeyFromSeed(result.seed_phrase)
            store.setViewKey(viewKeyResult.view_key_hex)
            store.setSpendingKey(viewKeyResult.spending_key_hex)

            const addressResult = await getAddressFromSeed(result.seed_phrase)
            store.setAddress(addressResult)

            return result
        } catch (err) {
            console.error('Error generating wallet:', err)
            throw err
        }
    }

    const getViewKeyFromSeed = async (seedPhrase: string[]): Promise<ViewKeyResult> => {
        const result = await (window as any).walletApi.generateKeysFromSeed([...seedPhrase])
        return JSON.parse(result)
    }

    const recoverWalletFromSeed = async (seedPhrase: string[]): Promise<WalletState> => {
        try {
            const isValid = validate_seed_phrase(JSON.stringify(seedPhrase))
            if (!isValid) throw new Error('Invalid seed phrase')

            const result = await getViewKeyFromSeed(seedPhrase)

            store.setSeedPhrase(seedPhrase)
            store.setReceiverId(result.receiver_identifier)
            store.setViewKey(result.view_key_hex)
            store.setSpendingKey(result.spending_key_hex)

            const addressResult = await getAddressFromSeed(seedPhrase)
            store.setAddress(addressResult)

            return {
                seedPhrase: seedPhrase,
                network: store.getNetwork,
                receiverId: result.receiver_identifier,
                viewKey: result.view_key_hex,
                spendingKey: result.spending_key_hex,
                address: addressResult,
            }
        } catch (err) {
            console.error('Error recovering wallet from seed:', err)
            throw err
        }
    }

    const getAddressFromSeed = async (seedPhrase: string[]): Promise<string> => {
        await ensureWasmInitialized()
        const seedPhraseJson = JSON.stringify(seedPhrase)
        return address_from_seed(seedPhraseJson, store.getNetwork)
    }

    const decryptKeystore = async (password: string): Promise<void> => {
        try {
            const keystorePath = store.getKeystorePath
            if (!keystorePath) await checkKeystore()

            const result = await (window as any).walletApi.decryptKeystore(
                store.getKeystorePath,
                password
            )

            const seedPhrase = result.phrase.trim().split(/\s+/)
            const viewKeyResult = await getViewKeyFromSeed(seedPhrase)

            store.setPassword(password)
            store.setSeedPhrase(seedPhrase)
            store.setViewKey(viewKeyResult.view_key_hex)
            store.setReceiverId(viewKeyResult.receiver_identifier)
            store.setSpendingKey(viewKeyResult.spending_key_hex)

            const addressResult = await getAddressFromSeed(seedPhrase)
            store.setAddress(addressResult)

            await loadMinBlockHeightFromKeystore()
        } catch (err) {
            if (
                err instanceof Error &&
                (err.message.includes('Unsupported state') ||
                    err.message.includes('unable to authenticate'))
            ) {
                console.error('Invalid password')
            } else console.error('Error decrypting keystore:', err)

            throw err
        }
    }

    const createKeystore = async (seed: string, password: string): Promise<string> => {
        try {
            const result = await (window as any).walletApi.createKeystore(seed, password)
            store.setKeystorePath(result.filePath)
            store.setMinBlockHeight(null)
            return result.filePath
        } catch (err) {
            console.error('Error creating keystore:', err)
            throw err
        }
    }

    const saveKeystoreAs = async (seed: string, password: string): Promise<string> => {
        try {
            const result = await (window as any).walletApi.saveKeystoreAs(seed, password)
            if (!result.filePath) throw new Error('User canceled')
            return result.filePath
        } catch (err) {
            console.error('Error saving keystore:', err)
            throw err
        }
    }

    const checkKeystore = async (): Promise<boolean> => {
        try {
            const keystoreFile = await (window as any).walletApi.checkKeystore()
            if (!keystoreFile.exists) return false

            store.setKeystorePath(keystoreFile.filePath)
            if ('minBlockHeight' in keystoreFile) {
                const height = keystoreFile.minBlockHeight
                store.setMinBlockHeight(
                    typeof height === 'number' && Number.isFinite(height) ? height : null
                )
            }
            return true
        } catch (err) {
            console.error('Error checking keystore:', err)
            throw err
        }
    }

    // ===== API METHODS =====

    const persistMinBlockHeight = async (utxos: any[]) => {
        const keystorePath = store.getKeystorePath

        if (!keystorePath) return

        try {
            const heights = utxos
                .map((utxo) => {
                    const rawHeight = [utxo?.blockHeight, utxo?.block_height, utxo?.height, utxo?.block?.height]
                        .find((h) => h !== null && h !== undefined) ?? null

                    const numericHeight =
                        typeof rawHeight === 'string' ? Number(rawHeight) : rawHeight

                    return Number.isFinite(numericHeight) ? Number(numericHeight) : null
                })
                .filter((height): height is number => height !== null)

            const minBlockHeight = heights.length ? Math.min(...heights) : null

            await (window as any).walletApi.updateMinBlockHeight(keystorePath, minBlockHeight)
            store.setMinBlockHeight(minBlockHeight)
        } catch (err) {
            console.error('Error saving min block height:', err)
        }
    }

    const loadMinBlockHeightFromKeystore = async (): Promise<number | null> => {
        const keystorePath = store.getKeystorePath

        if (!keystorePath) return null

        try {
            const response = await (window as any).walletApi.getMinBlockHeight(keystorePath)

            if (response?.success) {
                const minBlockHeight =
                    typeof response.minBlockHeight === 'number' && Number.isFinite(response.minBlockHeight)
                        ? response.minBlockHeight
                        : null

                store.setMinBlockHeight(minBlockHeight)
                return minBlockHeight
            }
        } catch (err) {
            console.error('Error loading min block height:', err)
        }

        return null
    }

    const getUtxos = async (): Promise<any> => {
        try {
            if (!store.getViewKey) {
                throw new Error('No view key available. Please import or generate a wallet first.')
            }

            let startBlock: number | null | undefined = store.getMinBlockHeight

            if (startBlock === null || startBlock === undefined) {
                startBlock = await loadMinBlockHeightFromKeystore()
            }

            const response = await API.getUtxosFromViewKey(
                store.getViewKey || '',
                typeof startBlock === 'number' && Number.isFinite(startBlock) ? startBlock : 0
            )

            const result = response?.result || response
            const utxoList = Array.isArray(result?.utxos)
                ? result.utxos
                : Array.isArray(result)
                ? result
                : []

            store.setUtxos(utxoList)

            await persistMinBlockHeight(utxoList)
            return result
        } catch (err) {
            console.error('Error getting UTXOs:', err)
            throw err
        }
    }

    const getBalance = async (): Promise<any> => {
        try {
            let startBlock: number | null | undefined = store.getMinBlockHeight

            if (startBlock === null || startBlock === undefined) {
                startBlock = await loadMinBlockHeightFromKeystore()
            }

            const response = await API.getBalance(
                store.getViewKey || '',
                typeof startBlock === 'number' && Number.isFinite(startBlock) ? startBlock : 0
            )
            const result = response?.result || response
            store.setBalance(result?.balance || result)
            store.setPendingBalance(result?.pendingBalance || result)
            return {
                balance: result?.balance || result,
                pendingBalance: result?.pendingBalance || result,
            }
        } catch (err) {
            console.error('Error getting balance:', err)
            throw err
        }
    }

    const getBlockHeight = async (): Promise<any> => {
        try {
            const response = await API.getBlockHeight()
            const result = response?.result || response
            return result?.height || result
        } catch (err) {
            console.error('Error getting block height:', err)
            throw err
        }
    }

    const getNetworkInfo = async (): Promise<any> => {
        try {
            const response = await API.getNetworkInfo()
            const result = response?.result || response
            store.setNetwork((result.network + 'net') as 'mainnet' | 'testnet')

            return result
        } catch (err) {
            console.error('Error getting network info:', err)
            throw err
        }
    }

    const buildTransaction = async (
        args: PayloadBuildTransaction
    ): Promise<any> => {
        let minBlockHeight: number | null | undefined = store.getMinBlockHeight
        if (minBlockHeight === null || minBlockHeight === undefined) {
            minBlockHeight = await loadMinBlockHeightFromKeystore()
        }

        const payload = {
            spendingKeyHex: store.getSpendingKey,
            inputAdditionRecords: args.inputAdditionRecords,
            minBlockHeight:
                typeof minBlockHeight === 'number' && Number.isFinite(minBlockHeight)
                    ? minBlockHeight
                    : 0,
            outputAddresses: args.outputAddresses,
            outputAmounts: args.outputAmounts,
            fee: args.fee,
        }
        return await (window as any).walletApi.buildTransaction(payload)
    }

    const broadcastSignedTransaction = async (transactionHex: string): Promise<any> => {
        try {
            const response = await API.broadcastSignedTransaction(transactionHex)
            const result = response?.result || response
            return result
        } catch (err) {
            console.error('Error sending transaction:', err)
            throw err
        }
    }

    const setNetwork = async (network: 'mainnet' | 'testnet') => {
        try {
            store.setNetwork(network)

            if (store.getSeedPhrase) {
                const viewKeyResult = await getViewKeyFromSeed(store.getSeedPhrase)
                store.setViewKey(viewKeyResult.view_key_hex)
                store.setSpendingKey(viewKeyResult.spending_key_hex)

                const addressResult = await getAddressFromSeed(store.getSeedPhrase)
                store.setAddress(addressResult)
            }
        } catch (err) {
            console.error('Error setting network:', err)
            throw err
        }
    }

    // ===== UTILITY METHODS =====
    const clearWallet = () => {
        store.clearWallet()
    }

    return {
        initWasm: ensureWasmInitialized,
        generateWallet,
        recoverWalletFromSeed,
        getViewKeyFromSeed,
        getAddressFromSeed,

        getUtxos,
        getBalance,
        getBlockHeight,
        getNetworkInfo,
        buildTransaction,
        broadcastSignedTransaction,
        decryptKeystore,
        createKeystore,
        saveKeystoreAs,
        checkKeystore,

        clearWallet,
        setNetwork,
    }
}