#!/usr/bin/env python3
# Copyright (c) 2022 The Bitcoin Unlimited developers
import asyncio
from test_framework.utilcompat import lock, unlock
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])


def new_tx(prevtx, out_script: CScript, amount, in_script: CScript):
    tx = create_transaction(
        prevtx=prevtx, value=amount, n=0, sig=unlock(in_script), out=lock(out_script)
    )
    pad_tx(tx)
    return tx


class ElectrumTokenGetVolumeTests(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_volume(n, cli, [coinbases.pop(0), coinbases.pop(0)])

        finally:
            cli.disconnect()

    async def test_volume(self, n, cli, coinbases):

        ##
        self.info("Test that volume shows up as unconfirmed, then confirmed.")
        tx1 = new_tx(
            coinbases.pop(0), out_script=SCRIPT1, amount=10000, in_script=CScript([])
        )
        tx2 = new_tx(
            coinbases.pop(0), out_script=SCRIPT1, amount=10000, in_script=CScript([])
        )
        n.sendrawtransaction(to_hex(tx1), True)
        n.sendrawtransaction(to_hex(tx2), True)
        await self.sync_mempool(cli, n)
        b = await cli.call(
            "blockchain.scripthash.get_volume", script_to_scripthash(lock(SCRIPT1))
        )
        assert_equal(
            {
                "confirmed_sats_in": 0,
                "confirmed_sats_out": 0,
                "mempool_sats_in": 20000,
                "mempool_sats_out": 0,
            },
            b,
        )
        await self.mine_blocks(cli, n, 1, [tx1, tx2])
        b = await cli.call(
            "blockchain.scripthash.get_volume", script_to_scripthash(lock(SCRIPT1))
        )
        assert_equal(
            {
                "confirmed_sats_in": 20000,
                "confirmed_sats_out": 0,
                "mempool_sats_in": 0,
                "mempool_sats_out": 0,
            },
            b,
        )

        ##
        self.info("Test spending a confirmed utxo shows up as out in mempool")
        tx3 = new_tx(tx2, out_script=SCRIPT2, amount=1000, in_script=SCRIPT1)
        n.sendrawtransaction(to_hex(tx3), True)
        await self.sync_mempool(cli, n)
        b = await cli.call(
            "blockchain.scripthash.get_volume", script_to_scripthash(lock(SCRIPT1))
        )
        assert_equal(
            {
                "confirmed_sats_in": 20000,
                "confirmed_sats_out": 0,
                "mempool_sats_in": 0,
                "mempool_sats_out": 10000,
            },
            b,
        )

        self.info("Check that out shows in confirmed")
        await self.mine_blocks(cli, n, 1, [tx3])
        b = await cli.call(
            "blockchain.scripthash.get_volume", script_to_scripthash(lock(SCRIPT1))
        )
        assert_equal(
            {
                "confirmed_sats_in": 20000,
                "confirmed_sats_out": 10000,
                "mempool_sats_in": 0,
                "mempool_sats_out": 0,
            },
            b,
        )

        ##
        self.info("Test receive and spend from same address in mempool")
        tx4 = new_tx(tx3, out_script=SCRIPT3, amount=400, in_script=SCRIPT2)
        tx5 = new_tx(tx4, out_script=SCRIPT3, amount=200, in_script=SCRIPT3)
        n.sendrawtransaction(to_hex(tx4), True)
        n.sendrawtransaction(to_hex(tx5), True)

        await self.sync_mempool(cli, n)
        b = await cli.call(
            "blockchain.scripthash.get_volume", script_to_scripthash(lock(SCRIPT3))
        )
        assert_equal(
            {
                "confirmed_sats_in": 0,
                "confirmed_sats_out": 0,
                "mempool_sats_in": 600,
                "mempool_sats_out": 400,
            },
            b,
        )


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