#!/usr/bin/env python3
"""
Tests for blockchain.address.get_first_spend and blockchain.scripthash.get_first_spend.
"""

import asyncio
from test_framework.util import assert_equal
from test_framework.electrumutil import ElectrumTestFramework, script_to_scripthash
from test_framework.electrumconnection import ElectrumConnection
from test_framework.constants import COIN
from test_framework.environment import on_nex, on_bch
from test_framework.script import CScript, OP_NOP, OP_FALSE, OP_DROP
from test_framework.utilcompat import lock, unlock
from test_framework.utiltx import pad_tx
from test_framework.serialize import to_hex

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


class ElectrumBlockchainAddressGetFirstSpend(ElectrumTestFramework):
    """Tests for get_first_spend RPC methods."""

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

        await self.bootstrap_p2p()
        cli = ElectrumConnection()
        await cli.connect()
        try:
            # Mine blocks using P2P to get mature coinbases
            coinbases = await self.mine_blocks(cli, n, 100)
            # Generate more blocks to mature the coinbase outputs for spending
            n.generate(100)

            await self.test_no_spends(cli)
            await self.test_spending_in_mempool(cli, n, coinbases)
            await self.test_spending_confirmed(cli, n, coinbases)
            await self.test_first_spend_after_multiple_spends(cli, n, coinbases)
            await self.test_comparison_with_get_first_use(cli, n, coinbases)
            await self.test_mempool_does_not_override_confirmed(cli, n, coinbases)
        finally:
            cli.disconnect()

    async def test_no_spends(self, cli):
        """Fresh scripthash with no spends should return null."""
        scriptpubkey = CScript([OP_NOP])
        scripthash = script_to_scripthash(lock(scriptpubkey))

        res = await cli.call("blockchain.scripthash.get_first_spend", scripthash)
        assert_equal(None, res)

    async def test_spending_in_mempool(self, cli, n, coinbases):
        """Test get_first_spend returns spend in mempool."""
        scriptpubkey = CScript([OP_NOP])
        scripthash = script_to_scripthash(lock(scriptpubkey))

        coinbase = coinbases.pop(0)
        fee = 200
        min_value = 10 * COIN
        tx_value = max(coinbase.vout[0].n_value - fee, min_value)

        # Create funding transaction: coinbase -> scriptpubkey
        fund_tx = self.create_transaction(
            coinbase,
            n=0,
            value=tx_value,
            sig=anyonecanspend_unlocking_script(),
            out=lock(scriptpubkey),
        )
        pad_tx(fund_tx)
        fund_txid = fund_tx.get_rpc_hex_id()
        await cli.call("blockchain.transaction.broadcast", to_hex(fund_tx))
        await self.wait_for_mempool_tx(cli, fund_txid, n)

        # Create spending transaction (in mempool)
        spend_tx = self.create_transaction(
            fund_tx,
            n=0,
            value=fund_tx.vout[0].n_value - fee,
            sig=unlock(scriptpubkey),
            out=lock(scriptpubkey),
        )
        pad_tx(spend_tx)
        spend_txid = spend_tx.get_rpc_hex_id()
        await cli.call("blockchain.transaction.broadcast", to_hex(spend_tx))
        await self.wait_for_mempool_tx(cli, spend_txid, n)

        # Verify get_first_spend returns mempool spend
        res = await cli.call("blockchain.scripthash.get_first_spend", scripthash)
        assert_equal(0, res["height"])
        assert_equal(
            "0000000000000000000000000000000000000000000000000000000000000000",
            res["block_hash"],
        )
        assert_equal(spend_txid, res["tx_hash"])

    async def test_spending_confirmed(self, cli, n, coinbases):
        """Test get_first_spend after confirmation."""
        scriptpubkey = CScript([OP_NOP, OP_NOP])
        scripthash = script_to_scripthash(lock(scriptpubkey))

        coinbase = coinbases.pop(0)
        fee = 200
        min_value = 10 * COIN
        tx_value = max(coinbase.vout[0].n_value - fee, min_value)

        # Fund
        fund_tx = self.create_transaction(
            coinbase,
            n=0,
            value=tx_value,
            sig=anyonecanspend_unlocking_script(),
            out=lock(scriptpubkey),
        )
        pad_tx(fund_tx)
        fund_txid = fund_tx.get_rpc_hex_id()
        await cli.call("blockchain.transaction.broadcast", to_hex(fund_tx))
        await self.wait_for_mempool_tx(cli, fund_txid, n)

        # Mine to confirm
        await self.mine_blocks(cli, n, 1)
        await self.sync_height(cli)

        # Create and broadcast spend
        spend_tx = self.create_transaction(
            fund_tx,
            n=0,
            value=fund_tx.vout[0].n_value - fee,
            sig=unlock(scriptpubkey),
            out=lock(scriptpubkey),
        )
        pad_tx(spend_tx)
        spend_txid = spend_tx.get_rpc_hex_id()
        await cli.call("blockchain.transaction.broadcast", to_hex(spend_tx))
        await self.wait_for_mempool_tx(cli, spend_txid, n)

        # Mine to confirm spend (use n.generate to pick up mempool txs)
        n.generate(1)
        await self.sync_height(cli)

        # Verify confirmed result
        res = await cli.call("blockchain.scripthash.get_first_spend", scripthash)
        assert_equal(n.getblockcount(), res["height"])
        assert_equal(n.getbestblockhash(), res["block_hash"])
        assert_equal(spend_txid, res["tx_hash"])

    async def test_first_spend_after_multiple_spends(self, cli, n, coinbases):
        # pylint: disable=too-many-locals
        """Test first_spend returns earliest spend, not most recent."""
        scriptpubkey = CScript([OP_NOP, OP_NOP, OP_NOP])
        scripthash = script_to_scripthash(lock(scriptpubkey))

        fee = 200
        min_value = 10 * COIN

        # Fund multiple times
        fund_txs = []
        for _ in range(3):
            coinbase = coinbases.pop(0)
            tx_value = max(coinbase.vout[0].n_value - fee, min_value)
            fund_tx = self.create_transaction(
                coinbase,
                n=0,
                value=tx_value,
                sig=anyonecanspend_unlocking_script(),
                out=lock(scriptpubkey),
            )
            pad_tx(fund_tx)
            fund_txs.append(fund_tx)
            await cli.call("blockchain.transaction.broadcast", to_hex(fund_tx))

        for fund_tx in fund_txs:
            await self.wait_for_mempool_tx(cli, fund_tx.get_rpc_hex_id(), n)

        # Mine to confirm all funding transactions
        await self.mine_blocks(cli, n, 1)
        await self.sync_height(cli)

        # Spend each UTXO
        first_spend_txid = None
        for fund_tx in fund_txs:
            spend_tx = self.create_transaction(
                fund_tx,
                n=0,
                value=fund_tx.vout[0].n_value - fee,
                sig=unlock(scriptpubkey),
                out=lock(scriptpubkey),
            )
            pad_tx(spend_tx)
            spend_txid = spend_tx.get_rpc_hex_id()
            if first_spend_txid is None:
                first_spend_txid = spend_txid
            await cli.call("blockchain.transaction.broadcast", to_hex(spend_tx))
            await self.wait_for_mempool_tx(cli, spend_txid, n)
            n.generate(1)
            await self.sync_height(cli)

        # Get first spend - should be earliest
        res = await cli.call("blockchain.scripthash.get_first_spend", scripthash)
        assert_equal(first_spend_txid, res["tx_hash"])

    async def test_comparison_with_get_first_use(self, cli, n, coinbases):
        # pylint: disable=too-many-locals
        """Verify get_first_spend and get_first_use return different txids."""
        scriptpubkey = CScript([OP_FALSE, OP_DROP])
        scripthash = script_to_scripthash(lock(scriptpubkey))

        coinbase = coinbases.pop(0)
        fee = 200
        min_value = 10 * COIN
        tx_value = max(coinbase.vout[0].n_value - fee, min_value)

        # Fund
        fund_tx = self.create_transaction(
            coinbase,
            n=0,
            value=tx_value,
            sig=anyonecanspend_unlocking_script(),
            out=lock(scriptpubkey),
        )
        pad_tx(fund_tx)
        fund_txid = fund_tx.get_rpc_hex_id()
        await cli.call("blockchain.transaction.broadcast", to_hex(fund_tx))
        await self.wait_for_mempool_tx(cli, fund_txid, n)

        # Mine to confirm
        await self.mine_blocks(cli, n, 1)
        await self.sync_height(cli)

        # Get first_use (funding tx)
        first_use = await cli.call("blockchain.scripthash.get_first_use", scripthash)

        # Create spend
        spend_tx = self.create_transaction(
            fund_tx,
            n=0,
            value=fund_tx.vout[0].n_value - fee,
            sig=unlock(scriptpubkey),
            out=lock(scriptpubkey),
        )
        pad_tx(spend_tx)
        spend_txid = spend_tx.get_rpc_hex_id()
        await cli.call("blockchain.transaction.broadcast", to_hex(spend_tx))
        await self.wait_for_mempool_tx(cli, spend_txid, n)
        n.generate(1)
        await self.sync_height(cli)

        # Get first_spend (spending tx)
        first_spend = await cli.call(
            "blockchain.scripthash.get_first_spend", scripthash
        )

        # They should be different
        assert_equal(fund_txid, first_use["tx_hash"])
        assert_equal(spend_txid, first_spend["tx_hash"])
        assert first_use["tx_hash"] != first_spend["tx_hash"]

    async def test_mempool_does_not_override_confirmed(self, cli, n, coinbases):
        # pylint: disable=too-many-locals
        """Test that mempool spend doesn't override confirmed first spend."""
        scriptpubkey = CScript([OP_NOP, OP_NOP, OP_NOP, OP_NOP])
        scripthash = script_to_scripthash(lock(scriptpubkey))

        fee = 200
        min_value = 10 * COIN

        # Fund and confirm (block A)
        coinbase1 = coinbases.pop(0)
        tx_value = max(coinbase1.vout[0].n_value - fee, min_value)
        fund_tx1 = self.create_transaction(
            coinbase1,
            n=0,
            value=tx_value,
            sig=anyonecanspend_unlocking_script(),
            out=lock(scriptpubkey),
        )
        pad_tx(fund_tx1)
        await cli.call("blockchain.transaction.broadcast", to_hex(fund_tx1))
        await self.wait_for_mempool_tx(cli, fund_tx1.get_rpc_hex_id(), n)

        await self.mine_blocks(cli, n, 1)
        await self.sync_height(cli)

        # First spend in next block (block B)
        spend_tx_1 = self.create_transaction(
            fund_tx1,
            n=0,
            value=fund_tx1.vout[0].n_value - fee,
            sig=unlock(scriptpubkey),
            out=lock(scriptpubkey),
        )
        pad_tx(spend_tx_1)
        spend_txid_1 = spend_tx_1.get_rpc_hex_id()
        await cli.call("blockchain.transaction.broadcast", to_hex(spend_tx_1))
        await self.wait_for_mempool_tx(cli, spend_txid_1, n)

        # Mine block B to confirm first spend (use n.generate to pick up mempool txs)
        n.generate(1)
        await self.sync_height(cli)

        block_b_height = n.getblockcount()

        # Verify first spend is at block B
        res = await cli.call("blockchain.scripthash.get_first_spend", scripthash)
        assert_equal(block_b_height, res["height"])
        assert_equal(spend_txid_1, res["tx_hash"])

        # Now fund again and create second spend but DON'T mine (stays in mempool)
        coinbase2 = coinbases.pop(0)
        tx_value2 = max(coinbase2.vout[0].n_value - fee, min_value)
        fund_tx2 = self.create_transaction(
            coinbase2,
            n=0,
            value=tx_value2,
            sig=anyonecanspend_unlocking_script(),
            out=lock(scriptpubkey),
        )
        pad_tx(fund_tx2)
        await cli.call("blockchain.transaction.broadcast", to_hex(fund_tx2))
        await self.wait_for_mempool_tx(cli, fund_tx2.get_rpc_hex_id(), n)
        await self.mine_blocks(cli, n, 1)
        await self.sync_height(cli)

        # Create second spend but don't mine - stays in mempool
        spend_tx_2 = self.create_transaction(
            fund_tx2,
            n=0,
            value=fund_tx2.vout[0].n_value - fee,
            sig=unlock(scriptpubkey),
            out=lock(scriptpubkey),
        )
        pad_tx(spend_tx_2)
        await cli.call("blockchain.transaction.broadcast", to_hex(spend_tx_2))
        await self.wait_for_mempool_tx(cli, spend_tx_2.get_rpc_hex_id(), n)

        # First spend should STILL return block B, not the mempool tx
        res = await cli.call("blockchain.scripthash.get_first_spend", scripthash)
        assert_equal(block_b_height, res["height"])
        assert_equal(spend_txid_1, res["tx_hash"])


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