#!/usr/bin/env python3
# Copyright (c) 2024 The Bitcoin Unlimited developers

"""
This tests that rostrum can connect and communicate with the node via p2p protocol
"""

import asyncio
from test_framework.script import CScript, OP_NOP
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
from test_framework.utilcompat import lock

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

TX_BROADCAST = "blockchain.transaction.broadcast"


class FeatureP2PTests(ElectrumTestFramework):
    async def run_test(self):
        await self.bootstrap_p2p()
        cli = ElectrumConnection()
        await cli.connect()
        try:
            coinbases = await self.mine_blocks(cli, self.nodes[0], 101)
            await self.test_p2p_broadcast(cli, self.nodes[0], coinbases.pop(0))

        finally:
            cli.disconnect()

    async def test_p2p_broadcast(self, cli, n, utxo):
        scriptpubkey = CScript([OP_NOP])

        tx = create_transaction(
            utxo,
            n=0,
            sig=anyonecanspend_unlocking_script(),
            out=lock(scriptpubkey),
            value=utxo.vout[0].n_value - 1000,
        )
        pad_tx(tx)

        n = self.nodes[0]

        # Mempool is empty before broadcast
        await self.wait_for_mempool_count(cli, count=0)

        broadcast_opt = {"p2p_only": True}

        txid = await get_txid_from_idem(
            n, await cli.call(TX_BROADCAST, tx.to_hex(), broadcast_opt)
        )

        # If the node sees p2p broadcated tx; it means p2p connection is good.
        await self.wait_for_mempool_tx(cli, txid)


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