#!/usr/bin/env python3
# Copyright (c) 2026 The Bitcoin Unlimited developers
"""
Tests for UTXO consistency after reorg (issues #242, #243).

After a reorg, verifies that:
- get_balance returns correct balance (restored after undo, zero after re-confirm)
- listunspent correctly reflects restored UTXOs
- transaction.get verbose can resolve inputs for surviving transactions
- transaction.get verbose works on re-confirmed transactions after reorg
"""

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

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

SCRIPTPUBKEY_S1 = CScript([OP_TRUE, OP_DROP])
SCRIPTPUBKEY_S2 = CScript([OP_TRUE, OP_TRUE, OP_DROP, OP_DROP])
SCRIPTPUBKEY_S3 = CScript([OP_FALSE, OP_DROP])
SCRIPTHASH_S1 = script_to_scripthash(lock(SCRIPTPUBKEY_S1))
SCRIPTHASH_S2 = script_to_scripthash(lock(SCRIPTPUBKEY_S2))


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

        await self.bootstrap_p2p()
        cli = ElectrumConnection()
        await cli.connect()
        try:
            coinbases = await self.mine_blocks(cli, n, 100)
            n.generate(100)
            n.generate(110)
            await self.test_balance_after_reorg(cli, n, coinbases)
            await self.test_verbose_get_after_reorg(cli, n, coinbases)
            await self.test_verbose_get_after_reconfirm(cli, n, coinbases)
        finally:
            cli.disconnect()

    async def test_balance_after_reorg(self, cli, n, coinbases):
        """
        Issue #242: After a reorg undoes a block that spent an output,
        the balance for the original scripthash should be restored.
        After re-spending and confirming, the balance should go back to zero.

        Sequence:
          1. T1 pays to S1, confirm in block N
          2. T2 spends from S1 to S2, confirm in block N+1
          3. Verify S1 balance is 0 (fully spent)
          4. Invalidate block N+1 (reorg)
          5. Verify S1 balance is restored (T1 output is unspent again)
          6. Verify S2 balance is 0 (T2 was undone)
          7. Re-mine to confirm T2 again
          8. Verify S1 balance returns to 0
        """
        coinbase = coinbases.pop(0)
        await self.sync_height(cli)
        await self.wait_for_mempool_count(cli, count=0)

        fee = 200
        tx1_value = max(coinbase.vout[0].n_value - fee, 10 * COIN)
        tx1 = self.create_transaction(
            coinbase,
            n=0,
            value=tx1_value,
            sig=anyonecanspend_unlocking_script(),
            out=lock(SCRIPTPUBKEY_S1),
        )
        pad_tx(tx1)

        self.info("Step 1: Broadcast T1 and confirm in block N")
        await cli.call("blockchain.transaction.broadcast", to_hex(tx1))
        await self.wait_for_mempool_count(cli, count=1)
        n.generate(1)
        await self.sync_height(cli)
        await self.wait_for_mempool_count(cli, count=0)

        balance_s1 = await cli.call("blockchain.scripthash.get_balance", SCRIPTHASH_S1)
        assert_equal(
            tx1_value, balance_s1["confirmed"], "S1 should have T1 value after confirm"
        )

        self.info("Step 2: T2 spends from S1 to S2, confirm in block N+1")
        tx2_value = max(tx1.vout[0].n_value - fee, 10 * COIN)
        tx2 = self.create_transaction(
            tx1,
            n=0,
            value=tx2_value,
            sig=unlock(SCRIPTPUBKEY_S1),
            out=lock(SCRIPTPUBKEY_S2),
        )
        pad_tx(tx2)
        await cli.call("blockchain.transaction.broadcast", to_hex(tx2))
        await self.wait_for_mempool_count(cli, count=1)
        n.generate(1)
        block_n_plus_1 = n.getbestblockhash()
        await self.sync_height(cli)
        await self.wait_for_mempool_count(cli, count=0)

        self.info("Step 3: Verify S1 is fully spent (balance 0)")
        balance_s1 = await cli.call("blockchain.scripthash.get_balance", SCRIPTHASH_S1)
        assert_equal(
            0,
            balance_s1["confirmed"],
            "S1 confirmed balance should be 0 after T2 spends it",
        )

        balance_s2 = await cli.call("blockchain.scripthash.get_balance", SCRIPTHASH_S2)
        assert_equal(tx2_value, balance_s2["confirmed"], "S2 should have T2 value")

        self.info("Step 4: Invalidate block N+1 (reorg)")
        n.invalidateblock(block_n_plus_1)
        self.hash_at_height.clear()
        await self.sync_height(cli)
        await asyncio.sleep(1)

        self.info("Step 5: Verify S1 balance is restored after reorg")

        async def s1_balance_restored():
            b = await cli.call("blockchain.scripthash.get_balance", SCRIPTHASH_S1)
            return b["confirmed"] == tx1_value

        await wait_for(
            10, s1_balance_restored, "S1 balance should be restored after reorg"
        )

        self.info("Step 6: Verify S2 balance is 0 after reorg (T2 was undone)")
        balance_s2 = await cli.call("blockchain.scripthash.get_balance", SCRIPTHASH_S2)
        assert_equal(
            0,
            balance_s2["confirmed"],
            "S2 confirmed balance should be 0 after reorg undid T2",
        )

        self.info("Step 7: Re-mine to confirm T2 again and verify S1 goes back to 0")
        n.generate(1)
        await self.sync_height(cli)
        await self.wait_for_mempool_count(cli, count=0)

        async def s1_balance_zero():
            b = await cli.call("blockchain.scripthash.get_balance", SCRIPTHASH_S1)
            return b["confirmed"] == 0

        await wait_for(
            10, s1_balance_zero, "S1 balance should be 0 after re-confirming T2"
        )

        self.info("test_balance_after_reorg passed")

    async def test_verbose_get_after_reorg(self, cli, n, coinbases):
        """
        Issue #243: After a reorg, transaction.get verbose should still
        resolve inputs for transactions in surviving blocks.

        Sequence:
          1. T1 pays to S1, confirm in block N
          2. T2 spends from S1 to S2, confirm in block N+1
          3. Invalidate block N+1
          4. transaction.get verbose on T1 should succeed (T1 still confirmed)
          5. listunspent on S1 should show the restored UTXO
        """
        coinbase = coinbases.pop(0)
        await self.sync_height(cli)
        await self.wait_for_mempool_count(cli, count=0)

        fee = 200
        min_value = 10 * COIN
        tx1_value = max(coinbase.vout[0].n_value - fee, min_value)
        tx1 = self.create_transaction(
            coinbase,
            n=0,
            value=tx1_value,
            sig=anyonecanspend_unlocking_script(),
            out=lock(SCRIPTPUBKEY_S1),
        )
        pad_tx(tx1)
        tx1_id = tx1.get_rpc_hex_id()

        self.info("Step 1: Broadcast T1 and confirm")
        await cli.call("blockchain.transaction.broadcast", to_hex(tx1))
        await self.wait_for_mempool_count(cli, count=1)
        n.generate(1)
        await self.sync_height(cli)
        await self.wait_for_mempool_count(cli, count=0)

        self.info("Step 2: T2 spends T1 output, confirm in block N+1")
        tx2_value = max(tx1.vout[0].n_value - fee, 10 * COIN)
        tx2 = self.create_transaction(
            tx1,
            n=0,
            value=tx2_value,
            sig=unlock(SCRIPTPUBKEY_S1),
            out=lock(SCRIPTPUBKEY_S2),
        )
        pad_tx(tx2)
        await cli.call("blockchain.transaction.broadcast", to_hex(tx2))
        await self.wait_for_mempool_count(cli, count=1)
        n.generate(1)
        block_n_plus_1 = n.getbestblockhash()
        await self.sync_height(cli)
        await self.wait_for_mempool_count(cli, count=0)

        self.info("Step 3: Invalidate block N+1 (reorg)")
        n.invalidateblock(block_n_plus_1)
        self.hash_at_height.clear()
        await self.sync_height(cli)
        await asyncio.sleep(1)

        self.info("Step 4: Verbose transaction.get on T1 (still confirmed)")
        result = await cli.call("blockchain.transaction.get", tx1_id, True)
        assert result["txid"] == tx1_id, "verbose get should return T1"
        assert result["height"] is not None, "T1 should still be confirmed"
        assert "vin" in result, "verbose result should have vin"

        self.info("Step 5: Verify S1 confirmed balance is restored after reorg")
        # Note: T2 goes back to mempool after the reorg, so listunspent
        # correctly filters O1 as spent-in-mempool. We check confirmed
        # balance instead, which reflects the on-chain state.

        async def s1_balance_restored():
            b = await cli.call("blockchain.scripthash.get_balance", SCRIPTHASH_S1)
            return b["confirmed"] == tx1_value

        await wait_for(
            10,
            s1_balance_restored,
            "S1 confirmed balance should be restored after reorg",
        )

        # Cleanup: re-mine (T2 re-confirmed from mempool)
        n.generate(1)
        await self.sync_height(cli)
        await self.wait_for_mempool_count(cli, count=0)

        self.info("test_verbose_get_after_reorg passed")

    async def test_verbose_get_after_reconfirm(self, cli, n, coinbases):
        """
        Issue #243: After a reorg undoes a block and the transaction is
        re-confirmed in a new block, transaction.get verbose should still
        be able to resolve the inputs.

        The reorg undo may corrupt the OutputIndex for the funding
        transaction's output. When the spending transaction is re-confirmed,
        verbose get fails with "could not find funding tx for outpoint".

        Sequence:
          1. T1 pays to S3, confirm in block N
          2. T2 spends from S3 to S2, confirm in block N+1
          3. Invalidate block N+1 (reorg, T2 back to mempool)
          4. Mine new block (T2 re-confirmed from mempool)
          5. transaction.get verbose on T2 should resolve inputs
        """
        coinbase = coinbases.pop(0)
        await self.sync_height(cli)
        await self.wait_for_mempool_count(cli, count=0)

        fee = 200
        min_value = 10 * COIN
        tx1_value = max(coinbase.vout[0].n_value - fee, min_value)
        tx1 = self.create_transaction(
            coinbase,
            n=0,
            value=tx1_value,
            sig=anyonecanspend_unlocking_script(),
            out=lock(SCRIPTPUBKEY_S3),
        )
        pad_tx(tx1)

        self.info("Step 1: Broadcast T1, confirm in block N")
        await cli.call("blockchain.transaction.broadcast", to_hex(tx1))
        await self.wait_for_mempool_count(cli, count=1)
        n.generate(1)
        await self.sync_height(cli)
        await self.wait_for_mempool_count(cli, count=0)

        self.info("Step 2: T2 spends T1 output, confirm in block N+1")
        tx2_value = max(tx1.vout[0].n_value - fee, 10 * COIN)
        tx2 = self.create_transaction(
            tx1,
            n=0,
            value=tx2_value,
            sig=unlock(SCRIPTPUBKEY_S3),
            out=lock(SCRIPTPUBKEY_S2),
        )
        pad_tx(tx2)
        tx2_id = tx2.get_rpc_hex_id()
        await cli.call("blockchain.transaction.broadcast", to_hex(tx2))
        await self.wait_for_mempool_count(cli, count=1)
        n.generate(1)
        block_n_plus_1 = n.getbestblockhash()
        await self.sync_height(cli)
        await self.wait_for_mempool_count(cli, count=0)

        self.info("Step 3: Invalidate block N+1 (reorg, T2 back to mempool)")
        n.invalidateblock(block_n_plus_1)
        self.hash_at_height.clear()
        await self.sync_height(cli)
        await asyncio.sleep(1)

        self.info("Step 4: Mine new block (T2 re-confirmed from mempool)")
        n.generate(1)
        await self.sync_height(cli)
        await self.wait_for_mempool_count(cli, count=0)

        self.info("Step 5: Verbose transaction.get on T2 (re-confirmed)")
        result = await cli.call("blockchain.transaction.get", tx2_id, True)
        assert result["txid"] == tx2_id, "verbose get should return T2"
        assert result["height"] is not None, "T2 should be confirmed"
        assert "vin" in result, "verbose result should have vin"
        assert len(result["vin"]) > 0, "T2 should have inputs"

        self.info("test_verbose_get_after_reconfirm passed")


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