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)
        } 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)
            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)
            return true
        } catch (err) {
            console.error('Error checking keystore:', err)
            throw err
        }
    }

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

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

            const response = await API.getUtxosFromViewKey(store.getViewKey || '')

            const result = response?.result || response
            store.setUtxos(result.utxos || result || [])
            return result
        } catch (err) {
            console.error('Error getting UTXOs:', err)
            throw err
        }
    }

    const getBalance = async (): Promise<any> => {
        try {
            const response = await API.getBalance(store.getViewKey || '')
            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 buildTransactionWithPrimitiveProof = async (
        args: PayloadBuildTransaction
    ): Promise<any> => {
        const payload = {
            spendingKeyHex: store.getSpendingKey,
            inputAdditionRecords: args.inputAdditionRecords,
            outputAddresses: args.outputAddresses,
            outputAmounts: args.outputAmounts,
            fee: args.fee,
        }
        return await (window as any).walletApi.buildTransactionWithPrimitiveProof(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,
        buildTransactionWithPrimitiveProof,
        broadcastSignedTransaction,
        decryptKeystore,
        createKeystore,
        saveKeystoreAs,
        checkKeystore,

        clearWallet,
        setNetwork,
    }
}