zcash-android-wallet-sdk/sdk-lib/src/androidTest/java/cash/z/ecc/android/sdk/util/TestWallet.kt

package cash.z.ecc.android.sdk.util

import androidx.test.platform.app.InstrumentationRegistry
import cash.z.ecc.android.bip39.Mnemonics
import cash.z.ecc.android.bip39.toSeed
import cash.z.ecc.android.sdk.SdkSynchronizer
import cash.z.ecc.android.sdk.Synchronizer
import cash.z.ecc.android.sdk.internal.Twig
import cash.z.ecc.android.sdk.internal.deriveUnifiedSpendingKey
import cash.z.ecc.android.sdk.internal.jni.RustDerivationTool
import cash.z.ecc.android.sdk.model.Account
import cash.z.ecc.android.sdk.model.BlockHeight
import cash.z.ecc.android.sdk.model.Testnet
import cash.z.ecc.android.sdk.model.WalletBalance
import cash.z.ecc.android.sdk.model.Zatoshi
import cash.z.ecc.android.sdk.model.ZcashNetwork
import co.electriccoin.lightwallet.client.model.LightWalletEndpoint
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.DelicateCoroutinesApi
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.delay
import kotlinx.coroutines.flow.first
import kotlinx.coroutines.launch
import kotlinx.coroutines.newFixedThreadPoolContext
import kotlinx.coroutines.runBlocking
import java.util.concurrent.TimeoutException

/**
 * A simple wallet that connects to testnet for integration testing. The intention is that it is
 * easy to drive and nice to use.
 */
@OptIn(DelicateCoroutinesApi::class)
class TestWallet(
    val seedPhrase: String,
    val alias: String = "TestWallet",
    val network: ZcashNetwork = ZcashNetwork.Testnet,
    val endpoint: LightWalletEndpoint = LightWalletEndpoint.Testnet,
    startHeight: BlockHeight? = null
) {
    constructor(
        backup: Backups,
        network: ZcashNetwork = ZcashNetwork.Testnet,
        alias: String = "TestWallet"
    ) : this(
        backup.seedPhrase,
        network = network,
        startHeight = if (network == ZcashNetwork.Mainnet) backup.mainnetBirthday else backup.testnetBirthday,
        alias = alias
    )

    val walletScope = CoroutineScope(
        SupervisorJob() + newFixedThreadPoolContext(3, this.javaClass.simpleName)
    )

    // Although runBlocking isn't great, this usage is OK because this is only used within the
    // automated tests

    private val account = Account.DEFAULT
    private val context = InstrumentationRegistry.getInstrumentation().context
    private val seed: ByteArray = Mnemonics.MnemonicCode(seedPhrase).toSeed()
    private val spendingKey =
        runBlocking { RustDerivationTool.new().deriveUnifiedSpendingKey(seed, network = network, account) }
    val synchronizer: SdkSynchronizer = Synchronizer.newBlocking(
        context,
        network,
        alias,
        lightWalletEndpoint = endpoint,
        seed = seed,
        startHeight
    ) as SdkSynchronizer

    val available get() = synchronizer.saplingBalances.value?.available
    val unifiedAddress =
        runBlocking { synchronizer.getUnifiedAddress(account) }
    val transparentAddress =
        runBlocking { synchronizer.getTransparentAddress(account) }
    val birthdayHeight get() = synchronizer.latestBirthdayHeight
    val networkName get() = synchronizer.network.networkName

    suspend fun transparentBalance(): WalletBalance {
        synchronizer.refreshUtxos(account, synchronizer.latestBirthdayHeight)
        return synchronizer.getTransparentBalance(transparentAddress)
    }

    suspend fun sync(timeout: Long = -1): TestWallet {
        val killSwitch = walletScope.launch {
            if (timeout > 0) {
                delay(timeout)
                throw TimeoutException("Failed to sync wallet within ${timeout}ms")
            }
        }

        // block until synced
        synchronizer.status.first { it == Synchronizer.Status.SYNCED }
        killSwitch.cancel()
        return this
    }

    suspend fun send(
        address: String = transparentAddress,
        memo: String = "",
        amount: Zatoshi = Zatoshi(500L)
    ): TestWallet {
        synchronizer.sendToAddress(spendingKey, amount, address, memo)
        return this
    }

    suspend fun rewindToHeight(height: BlockHeight): TestWallet {
        synchronizer.rewindToNearestHeight(height, false)
        return this
    }

    suspend fun shieldFunds(): TestWallet {
        synchronizer.refreshUtxos(Account.DEFAULT, BlockHeight.new(ZcashNetwork.Mainnet, 935000)).let { count ->
            Twig.debug { "FOUND $count new UTXOs" }
        }

        synchronizer.getTransparentBalance(transparentAddress).let { walletBalance ->
            Twig.debug { "FOUND utxo balance of total: ${walletBalance.total}  available: ${walletBalance.available}" }

            if (walletBalance.available.value > 0L) {
                synchronizer.shieldFunds(spendingKey)
            }
        }

        return this
    }

    suspend fun join(timeout: Long? = null): TestWallet {
        // block until stopped
        Twig.debug { "Staying alive until synchronizer is stopped!" }
        if (timeout != null) {
            Twig.debug { "Scheduling a stop in ${timeout}ms" }
            walletScope.launch {
                delay(timeout)
                synchronizer.close()
            }
        }
        synchronizer.status.first { it == Synchronizer.Status.STOPPED }
        Twig.debug { "Stopped!" }
        return this
    }

    companion object {
    }

    enum class Backups(val seedPhrase: String, val testnetBirthday: BlockHeight, val mainnetBirthday: BlockHeight) {
        // TODO: [#902] Get the proper birthday values for test wallets
        // TODO: [#902] https://github.com/zcash/zcash-android-wallet-sdk/issues/902
        DEFAULT(
            "column rhythm acoustic gym cost fit keen maze fence seed mail medal shrimp tell relief clip cannon foster soldier shallow refuse lunar parrot banana",
            BlockHeight.new(
                ZcashNetwork.Testnet,
                1_355_928
            ),
            BlockHeight.new(ZcashNetwork.Mainnet, 1_000_000)
        ),
        SAMPLE_WALLET(
            "input frown warm senior anxiety abuse yard prefer churn reject people glimpse govern glory crumble swallow verb laptop switch trophy inform friend permit purpose",
            BlockHeight.new(
                ZcashNetwork.Testnet,
                1_330_190
            ),
            BlockHeight.new(ZcashNetwork.Mainnet, 1_000_000)
        ),
        DEV_WALLET(
            "still champion voice habit trend flight survey between bitter process artefact blind carbon truly provide dizzy crush flush breeze blouse charge solid fish spread",
            BlockHeight.new(
                ZcashNetwork.Testnet,
                1_000_000
            ),
            BlockHeight.new(ZcashNetwork.Mainnet, 991645)
        ),
        ALICE(
            "quantum whisper lion route fury lunar pelican image job client hundred sauce chimney barely life cliff spirit admit weekend message recipe trumpet impact kitten",
            BlockHeight.new(
                ZcashNetwork.Testnet,
                1_330_190
            ),
            BlockHeight.new(ZcashNetwork.Mainnet, 1_000_000)
        ),
        BOB(
            "canvas wine sugar acquire garment spy tongue odor hole cage year habit bullet make label human unit option top calm neutral try vocal arena",
            BlockHeight.new(
                ZcashNetwork.Testnet,
                1_330_190
            ),
            BlockHeight.new(ZcashNetwork.Mainnet, 1_000_000)
        )
    }
}