#!/usr/bin/env python3
# Copyright (c) 2024 The Bitcoin Unlimited developers
import asyncio
from test_framework.nex.nodemessages import TxOut
from test_framework.serialize import to_hex
from test_framework.util import assert_equal
from test_framework.electrumutil import (
    ElectrumTestFramework,
)
from test_framework.electrumconnection import ElectrumConnection
from test_framework.environment import on_bch, on_nex
from test_framework.utiltx import p2pkh

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

#
# Tests that we can handle a large number of UTXOs.
#


def make_conform_to_ctor(txs):
    for tx in txs:
        tx.rehash()
    return sorted(txs, key=lambda tx: tx.get_rpc_hex_id())


class ManyUtxosTests(ElectrumTestFramework):
    """
    Test blockchain.address.listunspent with a large number of UTXOs
    """

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

        # Bootstrap P2P network
        await self.bootstrap_p2p()

        cli = ElectrumConnection()
        await cli.connect()
        try:
            # Mine blocks to get UTXOs
            coinbases = await self.mine_blocks(cli, n, 201)
            await self.test_many_utxos(n, cli, coinbases)
        finally:
            cli.disconnect()

    async def test_many_utxos(self, n, cli, coinbases):
        """
        Test listunspent with a large number of UTXOs
        """
        # Create a new address
        addr = n.getnewaddress("p2pkh")

        # Create 100 transactions, each with 200 outputs
        txs = []
        for _ in range(100):
            # Create a transaction with 200 outputs to our address
            coinbase = coinbases.pop(0)
            tx = create_transaction(
                prevtx=coinbase,
                value=10000,
                n=0,
                sig=anyonecanspend_unlocking_script(),
                out=p2pkh(addr),
                out_type=TxOut.TYPE_SATOSCRIPT,
            )

            # Add 199 more outputs
            for _ in range(199):
                tx.vout.append(tx.vout[0])

            txs.append(tx)

        # Sort transactions by txid in little-endian hex format
        txs = make_conform_to_ctor(txs)

        # Mine all transactions in a single block
        if on_nex():
            # may not fit in a single block initially; just use mempool
            for tx in txs:
                n.sendrawtransaction(to_hex(tx), True)
            await self.wait_for_mempool_count(cli, n, count=len(txs))
        else:
            await self.mine_blocks(cli, n, 1, txs)

        # Get listunspent for our address
        unspent = await cli.call("blockchain.address.listunspent", addr)
        print(f"Number of UTXOs: {len(unspent)}")
        print(f"First UTXO: {unspent[0]}")

        # Verify we have 20,000 UTXOs (100 transactions * 200 outputs)
        assert_equal(len(unspent), 20000)


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