#!/usr/bin/env python3
# Copyright (c) 2025 The Limenka developers
# Distributed under the MIT software license, see the accompanying
# file COPYING or http://www.opensource.org/licenses/mit-license.php.
"""Test blockreconstructionextratxn and blockreconstructionextratxnsize options with compact blocks."""

from test_framework.blocktools import (
    COINBASE_MATURITY,
    NORMAL_GBT_REQUEST_PARAMS,
    create_block,
    add_witness_commitment,
)
from test_framework.messages import (
    CTxOut,
    HeaderAndShortIDs,
    MSG_BLOCK,
    msg_cmpctblock,
    msg_sendcmpct,
    msg_tx,
    tx_from_hex,
)
from test_framework.p2p import (
    P2PInterface,
    p2p_lock,
)
from test_framework.script import (
    CScript,
    OP_DROP,
    OP_TRUE,
    OP_RETURN,
)
from test_framework.script_util import (
    keys_to_multisig_script,
)
from test_framework.test_framework import LimenkaTestFramework
from test_framework.util import (
    assert_equal,
    softfork_active,
)
from decimal import Decimal
from test_framework.wallet import MiniWallet


# TestP2PConn: A peer we use to send messages to limenkad, and store responses.
class TestP2PConn(P2PInterface):
    def __init__(self):
        super().__init__()
        self.last_sendcmpct = []
        self.block_announced = False
        # Store the hashes of blocks we've seen announced.
        # This is for synchronizing the p2p message traffic,
        # so we can eg wait until a particular block is announced.
        self.announced_blockhashes = set()

    def on_sendcmpct(self, message):
        self.last_sendcmpct.append(message)

    def on_cmpctblock(self, message):
        self.block_announced = True
        self.last_message["cmpctblock"].header_and_shortids.header.calc_sha256()
        self.announced_blockhashes.add(self.last_message["cmpctblock"].header_and_shortids.header.sha256)

    def on_headers(self, message):
        self.block_announced = True
        for x in self.last_message["headers"].headers:
            x.calc_sha256()
            self.announced_blockhashes.add(x.sha256)

    def on_inv(self, message):
        for x in self.last_message["inv"].inv:
            if x.type == MSG_BLOCK:
                self.block_announced = True
                self.announced_blockhashes.add(x.hash)

    # Requires caller to hold p2p_lock
    def received_block_announcement(self):
        return self.block_announced

    def clear_block_announcement(self):
        with p2p_lock:
            self.block_announced = False
            self.last_message.pop("inv", None)
            self.last_message.pop("headers", None)
            self.last_message.pop("cmpctblock", None)

    def clear_getblocktxn(self):
        with p2p_lock:
            self.last_message.pop("getblocktxn", None)


class CompactBlocksBlockReconstructionLimitTest(LimenkaTestFramework):
    def set_test_params(self):
        self.setup_clean_chain = True
        self.num_nodes = 1
        self.extra_args = [[
            "-acceptnonstdtxn=0",
            "-incrementalrelayfee=0.00001",
            "-debug=net",
        ]]
        self.utxos = []

    def build_block_on_tip(self, node):
        """Build a block on top of the current tip."""
        block = create_block(tmpl=node.getblocktemplate(NORMAL_GBT_REQUEST_PARAMS))
        block.solve()
        return block

    def make_utxos(self):
        """Generate blocks to create UTXOs for the wallet."""
        self.generate(self.wallet, COINBASE_MATURITY + 1000)

    def restart_node_with_limit(self, *, memory_mb=None, count=None):
        """Restart node with specific size and/or count limits."""
        extra_args = self.extra_args[0] + [
            "-datacarriersize=83",
        ]

        if memory_mb is not None:
            self.log.info(f"Setting size limit: {memory_mb} MB")
            extra_args.append(f"-blockreconstructionextratxnsize={memory_mb}")

        if count is not None:
            self.log.info(f"Setting transaction count limit: {count}")
            extra_args.append(f"-blockreconstructionextratxn={count}")

        self.log.info(f"Restarting node with args: {extra_args}")
        self.restart_node(0, extra_args=extra_args)
        self.segwit_node = self.nodes[0].add_p2p_connection(TestP2PConn())
        self.segwit_node.send_and_ping(msg_sendcmpct(announce=True, version=2))

    def create_policy_rejected_tx(self, rejection_type="dust", target_size=None):
        """Create a transaction that will be rejected for policy reasons but added to extra pool."""

        if rejection_type == "dust":
            tx_info = self.wallet.create_self_transfer()
            dust_amount = 100
            dust_script = CScript([OP_TRUE])
            tx_info['tx'].vout.append(CTxOut(dust_amount, dust_script))
            tx_info['tx'].vout[0].nValue -= dust_amount

        elif rejection_type == "low_fee":
            tx_info = self.wallet.create_self_transfer(fee_rate=Decimal('0.00000100'))

        elif rejection_type == "op_return_size":
            tx_info = self.wallet.create_self_transfer()
            data = b'x' * 85
            tx_info['tx'].vout.append(CTxOut(0, CScript([OP_RETURN, data])))

        elif rejection_type == "nonstandard_script":
            tx_info = self.wallet.create_self_transfer()
            pubkeys = []
            for _ in range(4):
                pubkeys.append(bytes([0x02] + [0x00] * 32))
            multisig_script = keys_to_multisig_script(pubkeys, k=4)
            tx_info['tx'].vout.append(CTxOut(10000, multisig_script))
            tx_info['tx'].vout[0].nValue -= 10000

        else:
            raise ValueError(f"Unknown rejection type: {rejection_type}")

        # Add padding outputs to reach target size
        if target_size:
            # Estimate current transaction size
            tx_info['tx'].rehash()
            base_size = len(tx_info['tx'].serialize())

            if base_size < target_size:
                # Each padded output approximately 200 bytes
                bytes_per_output = 200
                num_outputs = (target_size - base_size) // bytes_per_output

                for _ in range(num_outputs):
                    padding_data = b'x' * 190
                    script = CScript([padding_data, OP_DROP, OP_TRUE])
                    tx_info['tx'].vout.append(CTxOut(100, script))
                    tx_info['tx'].vout[0].nValue -= 100

        tx_info['tx'].rehash()
        tx_info['hex'] = tx_info['tx'].serialize().hex()
        return tx_info

    def populate_extra_pool(self, num_txs, rejection_type="dust", target_size=None):
        """Populate the extra transaction pool using policy-rejected transactions."""
        rejected_txs = []

        for i in range(num_txs):
            tx_info = self.create_policy_rejected_tx(rejection_type, target_size=target_size)
            tx_obj = tx_from_hex(tx_info['hex'])
            self.segwit_node.send_message(msg_tx(tx_obj))
            rejected_txs.append(tx_info)

        self.segwit_node.sync_with_ping()

        return rejected_txs

    def send_compact_block(self, transactions, indices):
        """Send a compact block and check which transactions are requested for reconstruction."""
        node = self.nodes[0]

        # Build block
        block = self.build_block_on_tip(node)

        for i in indices:
            tx_obj = tx_from_hex(transactions[i]['hex'])
            block.vtx.append(tx_obj)

        # Add witness commitment for blocks with witness transactions
        add_witness_commitment(block)
        block.solve()

        # Send as compact block
        cmpct_block = HeaderAndShortIDs()
        cmpct_block.initialize_from_block(block, use_witness=True)
        self.segwit_node.send_and_ping(msg_cmpctblock(cmpct_block.to_p2p()))

        # Check if node requested missing transactions
        with p2p_lock:
            getblocktxn = self.segwit_node.last_message.get("getblocktxn")

        num_tx_requested = len(getblocktxn.block_txn_request.indexes) if getblocktxn else 0
        self.segwit_node.clear_getblocktxn()

        # Convert differential encoding to absolute indices (from BlockTransactionRequest)
        missing_indices = []
        if getblocktxn:
            absolute_block_indices = getblocktxn.block_txn_request.to_absolute()
            # Convert from block positions to transaction indices (subtract 1 for coinbase)
            missing_indices = [idx - 1 for idx in absolute_block_indices]

        return {
            "block": block,
            "getblocktxn": getblocktxn,
            "num_tx_requested": num_tx_requested,
            "missing_indices": missing_indices
        }

    # TEST: policy-rejected transactions

    def test_policy_rejection_types(self):
        """Test that each policy rejection type adds transactions to extra pool."""
        self.log.info("Testing policy rejection types for extra pool...")

        self.restart_node_with_limit(count=100)

        rejection_types = ["dust", "low_fee", "op_return_size", "nonstandard_script"]
        rejected_txs = []

        for rejection_type in rejection_types:
            self.log.info(f"Testing {rejection_type} rejection...")
            tx_info = self.create_policy_rejected_tx(rejection_type)

            tx_obj = tx_from_hex(tx_info['hex'])
            self.segwit_node.send_message(msg_tx(tx_obj))

            rejected_txs.append({
                'type': rejection_type,
                'tx_info': tx_info,
                'txid': tx_info['tx'].hash,
                'wtxid': tx_info['tx'].getwtxid()
            })

        self.segwit_node.sync_with_ping()

        mempool = self.nodes[0].getrawmempool()
        for rejected in rejected_txs:
            assert_equal(rejected['txid'] in mempool, False)
            self.log.info(f"✓ {rejected['type']} transaction rejected from mempool")

        indices = list(range(len(rejected_txs)))
        tx_list = [r['tx_info'] for r in rejected_txs]

        result = self.send_compact_block(tx_list, indices)

        assert_equal(result["missing_indices"], [])
        self.log.info("✓ All rejected transactions are available in extra pool")

    # TEST: blockreconstructionextratxn

    def test_extratxnpool_disabled(self):
        """Test that setting count to 0 disables the extra transaction pool."""
        self.log.info("Testing disabled extra transaction pool (0 capacity)...")

        self.restart_node_with_limit(count=0)
        buffersize = 5
        rejected_txs = self.populate_extra_pool(buffersize)

        indices = list(range(buffersize))
        result = self.send_compact_block(rejected_txs, indices)
        assert_equal(result["missing_indices"], indices)
        self.log.info(f"✓ All {buffersize} transactions are missing (extra txn pool disabled)")

    def test_extratxnpool_capacity_and_wraparound(self):
        """Test extra transaction pool capacity and wraparound behavior."""
        self.log.info("Testing extra transaction pool capacity (400 transactions)...")

        buffersize = 400
        self.restart_node_with_limit(count=buffersize)

        rejected_txs = self.populate_extra_pool(buffersize)

        indices = list(range(buffersize))
        result = self.send_compact_block(rejected_txs, indices)

        assert_equal(result["missing_indices"], [])
        self.log.info("✓ All rejected transactions are in the extra txn pool")

        # Test that adding a 401st transaction causes eviction
        self.log.info("Adding transaction to test eviction...")
        self.populate_extra_pool(1)

        # Check original transactions again - first one should be evicted
        result2 = self.send_compact_block(rejected_txs, indices)
        assert_equal(result2["missing_indices"], [0])
        self.log.info("✓ Transaction 0 was evicted as expected (wraparound)")

    def test_single_extratxnpool_capacity(self):
        """Test edge case of single capacity extra transaction pool."""
        self.log.info("Testing single capacity extra transaction pool...")

        self.restart_node_with_limit(count=1)
        tx_count = 5

        rejected_txs = self.populate_extra_pool(tx_count)

        indices = list(range(tx_count))
        result = self.send_compact_block(rejected_txs, indices)

        expected_missing = list(range(tx_count - 1))
        assert_equal(result["missing_indices"], expected_missing)

    def test_extratxn_invalid_parameters(self):
        """Test handling of invalid blockreconstructionextratxn values."""
        self.log.info("Testing invalid parameter values...")

        # Test negative value - should be clamped to 0 (disabled)
        self.log.info("Testing negative value (-1)...")
        self.restart_node_with_limit(count=-1)

        # Add a transaction and verify pool is disabled
        rejected_txs = self.populate_extra_pool(1)
        result = self.send_compact_block(rejected_txs, [0])
        assert_equal(result["missing_indices"], [0])
        self.log.info("✓ Negative value correctly treated as disabled")

     # TEST: blockreconstructionextratxnsize

    def test_extratxnsize_zero_limit(self):
        """Test extra transaction pool zero size limit prevents extra txn pool."""
        self.log.info("Testing extra transaction pool zero size limit prevents extra txn pool...")
        self.restart_node_with_limit(memory_mb=0)

        rejected_txs = self.populate_extra_pool(1)
        result = self.send_compact_block(rejected_txs, [0])

        # Should fail - no size limit for extra pool
        assert result["getblocktxn"] is not None, "Node should try to request when zero size limit"
        assert_equal(int(self.nodes[0].getbestblockhash(), 16), result["block"].hashPrevBlock)

    def test_extratxnsize_eviction(self):
        """Test extra transaction pool size limit eviction behavior."""
        self.log.info("Testing extra transaction pool size limit eviction behavior...")

        buffersize = 60

        # First, test with 1MB limit - should fail
        self.log.info(f"Step 1: Testing with 1MB limit for {buffersize} large transactions")
        self.restart_node_with_limit(memory_mb=1, count=buffersize)

        # Create 60 large transactions (~20KB each = ~1.2MB total)
        # This exceeds the 1MB limit
        self.log.info(f"Creating {buffersize} large transactions (~20KB each, ~1.2MB total)")
        rejected_txs = self.populate_extra_pool(buffersize, target_size=20000)

        indices = list(range(buffersize))
        result_small = self.send_compact_block(rejected_txs, indices)

        # Should have evictions - can't fit 1.2MB in 1MB limit
        assert len(result_small["missing_indices"]) > 0, "1MB limit should cause evictions for 1.2MB of transactions"
        evicted_count = len(result_small["missing_indices"])
        self.log.info(f"✓ 1MB limit caused {evicted_count} evictions (can't fit ~1.2MB of transactions)")

        # Now test with larger size limit to show it succeeds
        self.log.info(f"Step 2: Testing with 2MB limit for same {buffersize} large transactions")
        self.restart_node_with_limit(memory_mb=2, count=buffersize)

        rejected_txs = self.populate_extra_pool(buffersize, target_size=20000)

        result_large = self.send_compact_block(rejected_txs, indices)

        # Should have NO evictions with 2MB limit
        assert result_large["missing_indices"] == [], "2MB limit should store all transactions"
        self.log.info(f"✓ 2MB limit successfully stores all {buffersize} large transactions (~1.2MB)")

    def test_extratxnsize_boundary(self):
        """Test extra transaction pool at exact size limit boundary."""
        self.log.info("Testing extra transaction pool exact size limit boundary...")

        limit_mb = 1
        self.restart_node_with_limit(memory_mb=limit_mb)

        test_count = 100
        rejected_txs = self.populate_extra_pool(test_count, target_size=20000)

        indices = list(range(test_count))
        result = self.send_compact_block(rejected_txs, indices)

        # Find the boundary - how many fit vs how many were evicted
        num_evicted = len(result["missing_indices"])
        num_fit = test_count - num_evicted

        # Now restart and add exactly the number that fit
        self.restart_node_with_limit(memory_mb=limit_mb)
        rejected_txs = self.populate_extra_pool(num_fit, target_size=20000)

        # Verify all fit
        indices = list(range(num_fit))
        result = self.send_compact_block(rejected_txs, indices)
        assert result["missing_indices"] == [], f"Expected all {num_fit} transactions to fit at boundary"

        # Add one more transaction - should evict exactly one
        self.log.info("Adding one more transaction at the boundary...")
        self.populate_extra_pool(1, target_size=20000)

        # Check original transactions again
        result2 = self.send_compact_block(rejected_txs, indices)
        assert len(result2["missing_indices"]) == 1, "Expected exactly 1 eviction at boundary"
        assert result2["missing_indices"] == [0], "Expected oldest transaction (0) to be evicted"

        self.log.info("Size limit boundary behavior verified - one transaction evicted when limit exceeded")

    def test_extratxnsize_small_limit(self):
        """Test extra transaction pool with very small size limit (0.1 MB)."""
        self.log.info("Testing extra transaction pool with 0.1 MB size limit...")

        limit_mb = 0.1
        self.restart_node_with_limit(memory_mb=limit_mb)

        test_count = 100
        rejected_txs = self.populate_extra_pool(test_count, target_size=500)
        indices = list(range(test_count))
        result = self.send_compact_block(rejected_txs, indices)

        # Find the boundary - how many fit vs how many were evicted
        num_evicted = len(result["missing_indices"])
        num_fit = test_count - num_evicted

        # fill exact capacity
        self.restart_node_with_limit(memory_mb=limit_mb)
        rejected_txs = self.populate_extra_pool(num_fit, target_size=500)

        indices = list(range(num_fit))
        result = self.send_compact_block(rejected_txs, indices)
        assert result["missing_indices"] == [], f"Expected all {num_fit} transactions to fit at boundary"

         # Add one more transaction - should evict exactly one
        self.log.info("Adding one more transaction at the boundary...")
        self.populate_extra_pool(1, target_size=500)

        # Check original transactions again
        result = self.send_compact_block(rejected_txs, indices)
        assert len(result["missing_indices"]) == 1, "Expected exactly 1 eviction at boundary"
        assert result["missing_indices"] == [0], "Expected oldest transaction (0) to be evicted"

        self.log.info("Small size limit boundary behavior verified - one transaction evicted when limit exceeded")

    def test_extratxnsize_large_transaction_exceeds_limit(self):
        """Test that a transaction larger than the entire size limit is rejected."""
        self.log.info("Testing large transaction that exceeds size limit...")

        limit_mb = 0.5 # 0.5 MB/500KB
        self.restart_node_with_limit(memory_mb=limit_mb)

        # Create a 600KB transaction (larger than 500KB limit)
        # This should be rejected from mempool AND not stored in extra pool
        rejected_txs = self.populate_extra_pool(1, target_size=600000) # 600 KB

        # try block reconstruction
        result = self.send_compact_block(rejected_txs, [0])

        # Verify the transaction is NOT available (was not stored in extra pool)
        assert len(result["missing_indices"]) == 1, "Large transaction should not be stored in extra pool"

    def run_test(self):
        self.wallet = MiniWallet(self.nodes[0])

        # Setup the p2p connection
        self.segwit_node = self.nodes[0].add_p2p_connection(TestP2PConn())

        # Create UTXOs for testing
        self.make_utxos()

        # Ensure segwit is active
        assert softfork_active(self.nodes[0], "segwit")

        # Test policy rejection types first
        self.test_policy_rejection_types()

        # Extra Txn capacity tests
        self.test_extratxnpool_disabled()
        self.test_extratxnpool_capacity_and_wraparound()
        self.test_single_extratxnpool_capacity()
        self.test_extratxn_invalid_parameters()

        # Extra Txn size tests
        self.test_extratxnsize_zero_limit()
        self.test_extratxnsize_eviction()
        self.test_extratxnsize_boundary()
        self.test_extratxnsize_small_limit()
        self.test_extratxnsize_large_transaction_exceeds_limit()


if __name__ == '__main__':
    CompactBlocksBlockReconstructionLimitTest(__file__).main()