#!/usr/bin/env python3
# Copyright (c) 2022 The Bitcoin Unlimited developers
import asyncio
from test_framework.electrumutil import script_to_scripthash
from test_framework.util import assert_equal
from test_framework.electrumutil import (
    ElectrumTestFramework,
)
from test_framework.environment import on_bch, on_nex
from test_framework.script import (
    CScript,
    OP_DROP,
    OP_FALSE,
    OP_TRUE,
)
from test_framework.serialize import to_hex
from test_framework.utiltx import pad_tx
from test_framework.electrumconnection import ElectrumConnection

if on_bch():
    from test_framework.bch.blocktools import create_transaction
elif on_nex():
    from test_framework.nex.blocktools import create_transaction
else:
    raise NotImplementedError()

SCRIPT1 = CScript([OP_FALSE, OP_DROP])
SCRIPT2 = CScript([OP_TRUE, OP_DROP])
SCRIPT3 = CScript([OP_FALSE, OP_FALSE, OP_DROP, OP_DROP])
SCRIPT4 = CScript([OP_TRUE, OP_FALSE, OP_DROP, OP_DROP])


SCRIPT1_HASH = script_to_scripthash(SCRIPT1)
SCRIPT2_HASH = script_to_scripthash(SCRIPT2)
SCRIPT3_HASH = script_to_scripthash(SCRIPT3)
SCRIPT4_HASH = script_to_scripthash(SCRIPT4)


def new_tx(prevtx, out_script, amount):
    tx = create_transaction(
        prevtx=prevtx, value=amount, n=0, sig=CScript([OP_TRUE]), out=out_script
    )
    pad_tx(tx)
    return tx


class ElectrumTokenGetBalanceTests(ElectrumTestFramework):
    async def run_test(self):
        n = self.nodes[0]

        await self.bootstrap_p2p()
        cli = ElectrumConnection()
        await cli.connect()
        coinbases = await self.mine_blocks(cli, n, 104)
        try:
            await self.test_balance(n, cli, [coinbases.pop(0), coinbases.pop(0)])
            await self.test_zeroes_out(n, cli, coinbases.pop(0))
            await self.test_negative_balance(n, cli, coinbases.pop(0))

        finally:
            cli.disconnect()

    async def test_balance(self, n, cli, coinbases):
        """
        Test that balance shows up as confirmed, then unconfirmed.
        """
        tx1 = new_tx(coinbases.pop(0), SCRIPT1, 10000)
        tx2 = new_tx(coinbases.pop(0), SCRIPT1, 10000)
        n.sendrawtransaction(to_hex(tx1), True)
        n.sendrawtransaction(to_hex(tx2), True)
        await self.sync_mempool_count(cli, n)
        b = await cli.call("blockchain.scripthash.get_balance", SCRIPT1_HASH)
        assert_equal(0, b["confirmed"])
        assert_equal(20000, b["unconfirmed"])
        await self.mine_blocks(cli, n, 1, [tx1, tx2])
        b = await cli.call("blockchain.scripthash.get_balance", SCRIPT1_HASH)
        assert_equal(20000, b["confirmed"])
        assert_equal(0, b["unconfirmed"])

    async def test_zeroes_out(self, n, cli, coinbase):
        """
        Special case. Sending and receiving the same
        amount results in unconfirmed balance of 0.

        This applies to both mempool and confirmed, if as
        long as the funding tx and spending tx is in the same
        index.
        """

        # Mempool
        fee = 500
        tx1 = new_tx(coinbase, SCRIPT4, 10000)
        tx2 = new_tx(tx1, SCRIPT2, 10000 - fee)
        n.sendrawtransaction(to_hex(tx1), True)
        n.sendrawtransaction(to_hex(tx2), True)
        await self.sync_mempool_count(cli)
        b = await cli.call("blockchain.scripthash.get_balance", SCRIPT4_HASH)
        assert_equal(0, b["unconfirmed"])

        b = await cli.call("blockchain.scripthash.get_balance", SCRIPT2_HASH)
        assert_equal(10000 - fee, b["unconfirmed"])

        # Confirmed index
        await self.mine_blocks(cli, n, 1, [tx1, tx2])
        b = await cli.call("blockchain.scripthash.get_balance", SCRIPT4_HASH)
        assert_equal(0, b["confirmed"])

        b = await cli.call("blockchain.scripthash.get_balance", SCRIPT2_HASH)
        assert_equal(10000 - fee, b["confirmed"])

    async def test_negative_balance(self, n, cli, coinbase):
        """
        Spending coins tokens results in negative mempool
        balance.
        """
        fee = 500
        tx1 = new_tx(coinbase, SCRIPT3, 10000)
        await self.mine_blocks(cli, n, 1, [tx1])

        # Spend from SCRIPT3
        tx2 = new_tx(tx1, SCRIPT2, 10000 - fee)
        n.sendrawtransaction(to_hex(tx2), True)
        await self.sync_mempool_count(cli, n)

        b = await cli.call("blockchain.scripthash.get_balance", SCRIPT3_HASH)
        assert_equal(-10000, b["unconfirmed"])
        assert_equal(10000, b["confirmed"])


if __name__ == "__main__":
    asyncio.run(ElectrumTokenGetBalanceTests().main())