Wallet.kt

/*
 * Copyright 2021-2025 thunderbiscuit and contributors.
 * Use of this source code is governed by the Apache 2.0 license that can be found in the ./LICENSE file.
 */

package org.bitcoindevkit.devkitwallet.domain

import android.util.Log
import kotlinx.coroutines.runBlocking
import org.bitcoindevkit.Address
import org.bitcoindevkit.AddressInfo
import org.bitcoindevkit.Amount
import org.bitcoindevkit.CanonicalTx
import org.bitcoindevkit.ChainPosition
import org.bitcoindevkit.Descriptor
import org.bitcoindevkit.DescriptorSecretKey
import org.bitcoindevkit.FeeRate
import org.bitcoindevkit.KeychainKind
import org.bitcoindevkit.Mnemonic
import org.bitcoindevkit.Network
import org.bitcoindevkit.Persister
import org.bitcoindevkit.Psbt
import org.bitcoindevkit.Script
import org.bitcoindevkit.TxBuilder
import org.bitcoindevkit.Update
import org.bitcoindevkit.WordCount
import org.bitcoindevkit.devkitwallet.data.ActiveWalletScriptType
import org.bitcoindevkit.devkitwallet.data.ConfirmationBlock
import org.bitcoindevkit.devkitwallet.data.NewWalletConfig
import org.bitcoindevkit.devkitwallet.data.RecoverWalletConfig
import org.bitcoindevkit.devkitwallet.data.SingleWallet
import org.bitcoindevkit.devkitwallet.data.Timestamp
import org.bitcoindevkit.devkitwallet.data.TxDetails
import org.bitcoindevkit.devkitwallet.domain.DwLogger.LogLevel.ERROR
import org.bitcoindevkit.devkitwallet.domain.DwLogger.LogLevel.INFO
import org.bitcoindevkit.devkitwallet.domain.utils.intoDomain
import org.bitcoindevkit.devkitwallet.domain.utils.intoProto
import org.bitcoindevkit.devkitwallet.presentation.viewmodels.Recipient
import java.io.File
import java.util.UUID
import org.bitcoindevkit.Wallet as BdkWallet

private const val TAG = "Wallet"

class Wallet private constructor(
    private val wallet: BdkWallet,
    private val walletSecrets: WalletSecrets,
    private val persister: Persister,
    private var fullScanCompleted: Boolean,
    private val walletId: String,
    private val userPreferencesRepository: UserPreferencesRepository,
    blockchainClientsConfig: BlockchainClientsConfig,
) {
    private var currentBlockchainClient: BlockchainClient? = blockchainClientsConfig.getClient()

    fun getWalletSecrets(): WalletSecrets {
        return walletSecrets
    }

    fun createTransaction(recipientList: List<Recipient>, feeRate: FeeRate): Psbt {
        // technique 1 for adding a list of recipients to the TxBuilder
        // var txBuilder = TxBuilder()
        // for (recipient in recipientList) {
        //     txBuilder  = txBuilder.addRecipient(address = recipient.first, amount = recipient.second)
        // }
        // txBuilder = txBuilder.feeRate(satPerVbyte = fee_rate)

        // technique 2 for adding a list of recipients to the TxBuilder
        val txBuilder =
            recipientList.fold(TxBuilder()) { builder, recipient ->
                // val address = Address(recipient.address)
                val scriptPubKey: Script = Address(recipient.address, Network.TESTNET).scriptPubkey()
                builder.addRecipient(scriptPubKey, Amount.fromSat(recipient.amount))
            }
        // if (!opReturnMsg.isNullOrEmpty()) {
        //     txBuilder = txBuilder.addData(opReturnMsg.toByteArray(charset = Charsets.UTF_8).asUByteArray().toList())
        // }
        return txBuilder.feeRate(feeRate).finish(wallet)
    }

    // @OptIn(ExperimentalUnsignedTypes::class)
    // fun createSendAllTransaction(
    //     recipient: String,
    //     feeRate: Float,
    //     enableRBF: Boolean,
    //     opReturnMsg: String?
    // ): PartiallySignedTransaction {
    //     val scriptPubkey: Script = Address(recipient).scriptPubkey()
    //     var txBuilder = TxBuilder()
    //         .drainWallet()
    //         .drainTo(scriptPubkey)
    //         .feeRate(satPerVbyte = feeRate)
    //
    //     if (enableRBF) {
    //         txBuilder = txBuilder.enableRbf()
    //     }
    //     if (!opReturnMsg.isNullOrEmpty()) {
    //         txBuilder = txBuilder.addData(opReturnMsg.toByteArray(charset = Charsets.UTF_8).asUByteArray().toList())
    //     }
    //     return txBuilder.finish(wallet).psbt
    // }

    // fun createBumpFeeTransaction(txid: String, feeRate: Float): PartiallySignedTransaction {
    //     return BumpFeeTxBuilder(txid = txid, newFeeRate = feeRate)
    //         .enableRbf()
    //         .finish(wallet = wallet)
    // }

    fun sign(psbt: Psbt): Boolean {
        return wallet.sign(psbt)
    }

    fun broadcast(signedPsbt: Psbt): String {
        val tx = signedPsbt.extractTx()
        currentBlockchainClient?.broadcast(tx) ?: throw IllegalStateException(
            "Blockchain client not initialized"
        )
        return tx.computeTxid().toString()
    }

    fun sweep(wif: String, feeRate: FeeRate): Psbt {
        val shortWif = wif.take(8) + "..."
        DwLogger.log(INFO, "Sweep started for WIF $shortWif (length=${wif.length}) on ${wallet.network()}")
        DwLogger.log(INFO, "Esplora endpoint: ${getClientEndpoint()}")

        val candidates = listOf(
            "pkh($wif)",
            "wpkh($wif)",
            "tr($wif)",
            "sh(wpkh($wif))",
        )

        var bestWallet: BdkWallet? = null
        var maxBalance = 0UL
        val candidateErrors = mutableListOf<String>()

        val net = wallet.network()
        val tempFiles = mutableListOf<File>()

        for (descString in candidates) {
            val label = descString.substringBefore("(")
            try {
                val descriptor = Descriptor(descString, net)

                val tempDir = System.getProperty("java.io.tmpdir")
                val tempFile = File(
                    "$tempDir/temp-${UUID.randomUUID().toString().take(8)}.sqlite3",
                )
                tempFiles.add(tempFile)
                val tempConnection = Persister.newSqlite(tempFile.absolutePath)

                val tempBdkWallet = BdkWallet.createSingle(
                    descriptor = descriptor,
                    network = net,
                    persister = tempConnection,
                )

                DwLogger.log(INFO, "Sweep: scanning $label ...")
                val fullScanRequest = tempBdkWallet.startFullScan().build()
                val update = currentBlockchainClient?.fullScan(
                    fullScanRequest = fullScanRequest,
                    stopGap = 3u,
                ) ?: throw IllegalStateException("Blockchain client not initialized")
                tempBdkWallet.applyUpdate(update)

                val bdkBalance = tempBdkWallet.balance()
                val balance = bdkBalance.total.toSat()
                DwLogger.log(
                    INFO,
                    "Sweep: $label balance: " +
                        "confirmed=${bdkBalance.confirmed.toSat()} " +
                        "pending=${bdkBalance.untrustedPending.toSat()} " +
                        "total=$balance",
                )
                Log.i(TAG, "Balance for $label: $balance")

                if (balance > maxBalance) {
                    maxBalance = balance
                    bestWallet = tempBdkWallet
                }
            } catch (e: Exception) {
                val msg = "$label failed: ${e.javaClass.simpleName} — ${e.message}"
                Log.e(TAG, msg, e)
                DwLogger.log(ERROR, "Sweep: $msg")
                candidateErrors.add(msg)
            }
        }

        try {
            if (bestWallet == null || maxBalance == 0UL) {
                val detail = if (candidateErrors.isNotEmpty()) {
                    candidateErrors.joinToString("; ")
                } else {
                    "All script types returned 0 balance"
                }
                throw IllegalStateException(
                    "No sweepable funds found for this WIF ($detail)",
                )
            }

            DwLogger.log(INFO, "Sweep: found $maxBalance sats, building tx")
            val destinationAddress = getNewAddress().address
            val destinationScriptPubKey = destinationAddress.scriptPubkey()

            val txBuilder = TxBuilder()
                .drainWallet()
                .drainTo(destinationScriptPubKey)
                .feeRate(feeRate)

            val psbt = txBuilder.finish(bestWallet)
            val signed = bestWallet.sign(psbt)
            if (!signed) {
                throw IllegalStateException(
                    "Failed to sign sweep transaction with the temporary wallet",
                )
            }

            DwLogger.log(INFO, "Sweep: signed successfully")
            return psbt
        } finally {
            for (f in tempFiles) {
                listOf(f, File("${f.absolutePath}-wal"), File("${f.absolutePath}-shm"))
                    .forEach {
                        try {
                            it.delete()
                        } catch (_: Exception) {
                        }
                    }
            }
        }
    }

    private fun getAllTransactions(): List<CanonicalTx> = wallet.transactions()

    fun getAllTxDetails(): List<TxDetails> {
        val transactions = getAllTransactions()
        return transactions.map { tx ->
            val txid = tx.transaction.computeTxid()
            val (sent, received) = wallet.sentAndReceived(tx.transaction)
            var feeRate: FeeRate? = null
            var fee: Amount? = null
            // TODO: I don't know why we're getting negative fees here, but it looks like a bug
            try {
                fee = wallet.calculateFee(tx.transaction)
            } catch (e: Exception) {
                Log.e(TAG, "Error calculating fee rate for tx $txid: $e")
            }
            try {
                feeRate = wallet.calculateFeeRate(tx.transaction)
            } catch (e: Exception) {
                Log.e(TAG, "Error calculating fee for tx $txid: $e")
            }

            val (confirmationBlock, confirmationTimestamp, pending) =
                when (val position = tx.chainPosition) {
                    is ChainPosition.Unconfirmed -> {
                        Triple(null, null, true)
                    }

                    is ChainPosition.Confirmed -> {
                        Triple(
                            ConfirmationBlock(position.confirmationBlockTime.blockId.height),
                            Timestamp(position.confirmationBlockTime.confirmationTime),
                            false
                        )
                    }
                }
            TxDetails(
                tx.transaction,
                txid.toString(),
                sent.toSat(),
                received.toSat(),
                fee?.toSat() ?: 0uL,
                feeRate,
                pending,
                confirmationBlock,
                confirmationTimestamp
            )
        }
    }

    // fun getTransaction(txid: String): TransactionDetails? {
    //     val allTransactions = getAllTransactions()
    //     allTransactions.forEach {
    //         if (it.txid == txid) {
    //             return it
    //         }
    //     }
    //     return null
    // }

    private fun fullScan() {
        val fullScanRequest = wallet.startFullScan().build()
        val update: Update =
            currentBlockchainClient?.fullScan(
                fullScanRequest = fullScanRequest,
                stopGap = 20u,
            ) ?: throw IllegalStateException("Blockchain client not initialized")
        wallet.applyUpdate(update)
        wallet.persist(persister)
    }

    fun sync() {
        if (!fullScanCompleted) {
            Log.i(TAG, "Full scan required")
            fullScan()
            runBlocking {
                userPreferencesRepository.setFullScanCompleted(walletId)
                fullScanCompleted = true
            }
        } else {
            Log.i(TAG, "Just a normal sync!")
            val syncRequest = wallet.startSyncWithRevealedSpks().build()
            val update =
                currentBlockchainClient?.sync(
                    syncRequest = syncRequest,
                ) ?: throw IllegalStateException("Blockchain client not initialized")
            wallet.applyUpdate(update)
            wallet.persist(persister)
        }
    }

    fun getBalance(): ULong = wallet.balance().total.toSat()

    fun getNewAddress(): AddressInfo = wallet.revealNextAddress(KeychainKind.EXTERNAL)

    fun getClientEndpoint(): String = currentBlockchainClient?.endpoint() ?: "No active client"

    companion object {
        fun createWallet(
            newWalletConfig: NewWalletConfig,
            internalAppFilesPath: String,
            userPreferencesRepository: UserPreferencesRepository,
        ): Wallet {
            val mnemonic = Mnemonic(WordCount.WORDS12)
            val bip32ExtendedRootKey = DescriptorSecretKey(newWalletConfig.network, mnemonic, null)
            val descriptor: Descriptor = createScriptAppropriateDescriptor(
                newWalletConfig.scriptType,
                bip32ExtendedRootKey,
                newWalletConfig.network,
                KeychainKind.EXTERNAL
            )
            val changeDescriptor: Descriptor = createScriptAppropriateDescriptor(
                newWalletConfig.scriptType,
                bip32ExtendedRootKey,
                newWalletConfig.network,
                KeychainKind.INTERNAL
            )
            val walletId = UUID.randomUUID().toString()
            val persister = Persister.newSqlite("$internalAppFilesPath/wallet-${walletId.take(8)}.sqlite3")

            // Create SingleWallet object for saving to datastore
            val newWalletForDatastore: SingleWallet = SingleWallet
                .newBuilder()
                .setId(walletId)
                .setName(newWalletConfig.name)
                .setNetwork(newWalletConfig.network.intoProto())
                .setScriptType(newWalletConfig.scriptType)
                .setDescriptor(descriptor.toStringWithSecret())
                .setChangeDescriptor(changeDescriptor.toStringWithSecret())
                .setRecoveryPhrase(mnemonic.toString())
                .build()

            // TODO: launch this correctly, not on the main thread
            // Save the new wallet to the datastore
            runBlocking { userPreferencesRepository.updateActiveWallets(newWalletForDatastore) }

            val bdkWallet =
                BdkWallet(
                    descriptor = descriptor,
                    changeDescriptor = changeDescriptor,
                    network = newWalletConfig.network,
                    persister = persister,
                )

            val walletSecrets = WalletSecrets(descriptor, changeDescriptor, mnemonic.toString())

            return Wallet(
                wallet = bdkWallet,
                walletSecrets = walletSecrets,
                persister = persister,
                fullScanCompleted = false,
                walletId = walletId,
                userPreferencesRepository = userPreferencesRepository,
                blockchainClientsConfig = BlockchainClientsConfig.createDefaultConfig(newWalletConfig.network)
            )
        }

        fun loadActiveWallet(
            activeWallet: SingleWallet,
            internalAppFilesPath: String,
            userPreferencesRepository: UserPreferencesRepository,
        ): Wallet {
            val descriptor = Descriptor(activeWallet.descriptor, activeWallet.network.intoDomain())
            val changeDescriptor = Descriptor(activeWallet.changeDescriptor, activeWallet.network.intoDomain())
            val persister = Persister.newSqlite("$internalAppFilesPath/wallet-${activeWallet.id.take(8)}.sqlite3")
            val bdkWallet = BdkWallet.load(
                descriptor = descriptor,
                changeDescriptor = changeDescriptor,
                persister = persister
            )

            val walletSecrets = WalletSecrets(descriptor, changeDescriptor, activeWallet.recoveryPhrase)
            return Wallet(
                wallet = bdkWallet,
                walletSecrets = walletSecrets,
                persister = persister,
                fullScanCompleted = activeWallet.fullScanCompleted,
                walletId = activeWallet.id,
                userPreferencesRepository = userPreferencesRepository,
                blockchainClientsConfig = BlockchainClientsConfig.createDefaultConfig(activeWallet.network.intoDomain())
            )
        }

        fun recoverWallet(
            recoverWalletConfig: RecoverWalletConfig,
            internalAppFilesPath: String,
            userPreferencesRepository: UserPreferencesRepository,
        ): Wallet {
            Log.i(TAG, "Recovering wallet with config: $recoverWalletConfig")
            var descriptor: Descriptor? = null
            var changeDescriptor: Descriptor? = null
            var mnemonicString: String = ""

            // If there is a recovery phrase, we use it to recover the wallet
            if (recoverWalletConfig.recoveryPhrase != null && recoverWalletConfig.scriptType != null) {
                val mnemonic: Mnemonic = Mnemonic.fromString(recoverWalletConfig.recoveryPhrase)
                mnemonicString = mnemonic.toString()
                val bip32ExtendedRootKey = DescriptorSecretKey(recoverWalletConfig.network, mnemonic, null)
                descriptor = createScriptAppropriateDescriptor(
                    recoverWalletConfig.scriptType,
                    bip32ExtendedRootKey,
                    recoverWalletConfig.network,
                    KeychainKind.EXTERNAL
                )
                changeDescriptor = createScriptAppropriateDescriptor(
                    recoverWalletConfig.scriptType,
                    bip32ExtendedRootKey,
                    recoverWalletConfig.network,
                    KeychainKind.INTERNAL
                )
            } else {
                descriptor = recoverWalletConfig.descriptor
                changeDescriptor = recoverWalletConfig.changeDescriptor
            }
            val walletId = UUID.randomUUID().toString()
            val persister = Persister.newSqlite("$internalAppFilesPath/wallet-${walletId.take(8)}.sqlite3")

            // Create SingleWallet object for saving to datastore
            val newWalletForDatastore: SingleWallet = SingleWallet
                .newBuilder()
                .setId(walletId)
                .setName(recoverWalletConfig.name)
                .setNetwork(recoverWalletConfig.network.intoProto())
                .setScriptType(recoverWalletConfig.scriptType ?: ActiveWalletScriptType.UNKNOWN)
                .setDescriptor(descriptor.toStringWithSecret())
                .setChangeDescriptor(changeDescriptor.toStringWithSecret())
                .setRecoveryPhrase(mnemonicString)
                .build()

            // TODO: launch this correctly, not on the main thread
            // Save the new wallet to the datastore
            runBlocking { userPreferencesRepository.updateActiveWallets(newWalletForDatastore) }

            val bdkWallet = BdkWallet(
                descriptor = descriptor,
                changeDescriptor = changeDescriptor,
                persister = persister,
                network = recoverWalletConfig.network,
            )

            val walletSecrets = WalletSecrets(descriptor, changeDescriptor, mnemonicString)
            return Wallet(
                wallet = bdkWallet,
                walletSecrets = walletSecrets,
                persister = persister,
                fullScanCompleted = false,
                walletId = walletId,
                userPreferencesRepository = userPreferencesRepository,
                blockchainClientsConfig = BlockchainClientsConfig.createDefaultConfig(recoverWalletConfig.network)
            )
        }
    }
}

fun createScriptAppropriateDescriptor(
    scriptType: ActiveWalletScriptType,
    bip32ExtendedRootKey: DescriptorSecretKey,
    network: Network,
    keychain: KeychainKind,
): Descriptor {
    return if (keychain == KeychainKind.EXTERNAL) {
        when (scriptType) {
            ActiveWalletScriptType.P2WPKH -> Descriptor.newBip84(bip32ExtendedRootKey, KeychainKind.EXTERNAL, network)
            ActiveWalletScriptType.P2TR -> Descriptor.newBip86(bip32ExtendedRootKey, KeychainKind.EXTERNAL, network)
            ActiveWalletScriptType.UNKNOWN -> TODO()
            ActiveWalletScriptType.UNRECOGNIZED -> TODO()
        }
    } else {
        when (scriptType) {
            ActiveWalletScriptType.P2WPKH -> Descriptor.newBip84(bip32ExtendedRootKey, KeychainKind.INTERNAL, network)
            ActiveWalletScriptType.P2TR -> Descriptor.newBip86(bip32ExtendedRootKey, KeychainKind.INTERNAL, network)
            ActiveWalletScriptType.UNKNOWN -> TODO()
            ActiveWalletScriptType.UNRECOGNIZED -> TODO()
        }
    }
}

data class WalletSecrets(
    val descriptor: Descriptor,
    val changeDescriptor: Descriptor,
    val recoveryPhrase: String,
)