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

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()

GET_UTXO = "blockchain.utxo.get"
TX_BROADCAST = "blockchain.transaction.broadcast"


class ElectrumUtxoTests(ElectrumTestFramework):
    async def run_test(self):
        await self.bootstrap_p2p()
        cli = ElectrumConnection()
        await cli.connect()
        coinbases = await self.mine_blocks(cli, self.nodes[0], 104)

        try:
            await self.test_get_all(
                cli, self.nodes[0], [coinbases.pop(0), coinbases.pop(0)]
            )
            await self.test_scriptsig_filter(
                cli, self.nodes[0], [coinbases.pop(0), coinbases.pop(0)]
            )

        finally:
            cli.disconnect()

    async def test_get_all(self, cli, n, utxos):
        scriptpubkey = CScript([OP_NOP])

        tx1 = create_transaction(
            utxos[0],
            n=0,
            sig=CScript([OP_TRUE]),
            out=scriptpubkey,
            value=utxos[0].vout[0].n_value - 1000,
        )
        pad_tx(tx1)
        tx2 = create_transaction(
            utxos[1],
            n=0,
            sig=CScript([OP_TRUE]),
            out=scriptpubkey,
            value=utxos[1].vout[0].n_value - 1000,
        )
        pad_tx(tx2)

        n = self.nodes[0]

        # Mempool is empty before broadcast
        mempool = await cli.call("mempool.get")
        assert_equal(0, len(mempool["transactions"]))

        tx1_id = await get_txid_from_idem(n, await cli.call(TX_BROADCAST, tx1.to_hex()))
        tx2_id = await get_txid_from_idem(n, await cli.call(TX_BROADCAST, tx2.to_hex()))

        await self.wait_for_mempool_count(cli, count=2)
        mempool = await cli.call("mempool.get")
        print(mempool)
        assert_equal(2, len(mempool["transactions"]))
        assert tx1_id in mempool["transactions"]
        assert tx2_id in mempool["transactions"]

        # Clean mempool
        await self.mine_blocks(cli, n, 1, [tx1, tx2])

    async def test_scriptsig_filter(self, cli, n, utxos):
        scriptpubkey = CScript([OP_DROP, OP_TRUE])

        tx1 = create_transaction(
            utxos[0],
            n=0,
            sig=CScript([OP_TRUE]),
            out=scriptpubkey,
            value=utxos[0].vout[0].n_value - 1000,
        )
        pad_tx(tx1)
        tx2 = create_transaction(
            tx1,
            n=0,
            sig=CScript([b"FILTERME"]),
            out=scriptpubkey,
            value=tx1.vout[0].n_value - 1000,
        )
        pad_tx(tx2)

        n = self.nodes[0]

        # Mempool is empty before broadcast
        mempool = await cli.call("mempool.get")
        assert_equal(0, len(mempool["transactions"]))

        tx1_id = await get_txid_from_idem(n, await cli.call(TX_BROADCAST, tx1.to_hex()))
        tx2_id = await get_txid_from_idem(n, await cli.call(TX_BROADCAST, tx2.to_hex()))

        await self.wait_for_mempool_count(cli, count=2)
        mempool = await cli.call(
            "mempool.get", {"scriptsig": to_hex(b"FILTERME").decode("ascii")}
        )
        print(mempool)
        assert_equal(1, len(mempool["transactions"]))
        assert tx1_id not in mempool["transactions"]
        assert tx2_id in mempool["transactions"]

        # Clean mempool
        await self.mine_blocks(cli, n, 1, [tx1, tx2])


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