#!/usr/bin/env python3
# Copyright (c) 2023-2024 The Bitcoin developers
# Distributed under the MIT software license, see the accompanying
# file COPYING or http://www.opensource.org/licenses/mit-license.php.
"""Test support for ABLA EBAA which activates with Upgrade10."""
from test_framework.abla import AblaState
from test_framework.cdefs import (
    BLOCK_MAXBYTES_MAXSIGCHECKS_RATIO,
    DEFAULT_CONSENSUS_BLOCK_SIZE,
)
from test_framework.test_framework import BitcoinTestFramework
from test_framework.util import (
    assert_equal,
    assert_greater_than,
    assert_greater_than_or_equal,
    assert_not_equal,
    connect_nodes,
    disconnect_nodes,
)


class AblaTest(BitcoinTestFramework):

    def set_test_params(self):
        self.num_nodes = 2
        self.setup_clean_chain = True
        self.base_extra_args = ['-expire=0', '-checkmempool=0', '-allowunconnectedmining=1']
        self.extra_args = [['-upgrade10activationheight=2147483647', '-percentblockmaxsize=100']
                           + self.base_extra_args,
                           ['-upgrade10activationheight=0']]
        # We need a long rpc timeout here so that the sanitizer-undefined CI job may pass
        self.rpc_timeout = 600

    def get_header(self, blockhash=None, verbose=True):
        if blockhash is None:
            blockhash = self.nodes[0].getblockchaininfo()['bestblockhash']
        return self.nodes[0].getblockheader(blockhash, verbose)

    def get_height(self, blockhash=None) -> int:
        return self.get_header(blockhash, True)['height']

    def get_abla_activation_height(self):
        """Walk backwards from tip, finding ABLA activation height. Returns the height or None if ABLA not activated."""
        node = self.nodes[0]
        seen = False
        bh = node.getblockchaininfo()['bestblockhash']
        while bh:
            hdr = self.get_header(bh)
            if 'ablastate' in hdr:
                seen = True
            else:
                if seen:
                    return hdr['height'] + 1
                else:
                    return None
            bh = hdr.get('previousblockhash')
        return 0

    def run_test(self):
        node = self.nodes[0]
        addr, privkey = node.get_deterministic_priv_key()

        # We operate on node 0 disconnected from 1, and connect them at the end to verify that node1 was able to IBD
        # to node 0 correctly.
        disconnect_nodes(self.nodes[0], self.nodes[1])

        # Start off 120 MB mempool (~= about 50MB serialized size)
        node.fillmempool(120)
        assert_greater_than(node.getmempoolinfo()['bytes'], DEFAULT_CONSENSUS_BLOCK_SIZE)
        assert_greater_than_or_equal(node.getblockchaininfo()['blocks'], 100)

        # Ensure no ABLA state in headers
        assert 'ablastate' not in self.get_header()

        # Check that blocksize limit is at the consensus default
        assert_equal(node.getblocktemplatelight({})['sizelimit'], DEFAULT_CONSENSUS_BLOCK_SIZE)

        saved_height = self.get_height()

        # Now, enable the upgrade
        self.restart_node(0, extra_args=self.base_extra_args + [f'-upgrade10activationheight={saved_height + 1}',
                                                                '-percentblockmaxsize=100'])

        # No ABLA yet
        assert_equal(self.get_abla_activation_height(), None)

        # Mine 2 blocks to ensure activation
        node.generatetoaddress(2, addr)

        activation_height = self.get_abla_activation_height()
        assert_not_equal(activation_height, None)
        assert_equal(activation_height, saved_height + 1)

        node.fillmempool(290)  # Put 290 MB unserialized =~ 70 MB serialized tx data into mempool

        abla = AblaState.FromRpcDict(self.get_header())  # Get tip's abla state
        new_block_hash = node.generatetoaddress(1, addr)[0]  # Mine a full block
        new_block_size = node.getblock(new_block_hash, 1)['size']
        self.log.info(f"Mined block size: {new_block_size}")
        assert_greater_than(new_block_size, abla.GetNextBlockSizeLimit() * 0.90)  # Ensure block was indeed full-ish

        new_abla_state = abla.NextBlockState(new_block_size)  # Advance abla state

        # Should match what we expect
        assert_equal(new_abla_state, AblaState.FromRpcDict(self.get_header()))

        # Since last block was full, limit should have grown
        next_block_size_limit = new_abla_state.GetNextBlockSizeLimit()
        assert_greater_than(next_block_size_limit, abla.GetNextBlockSizeLimit())

        # Ensure mining (using percentmaxblocksize=100.0) reports correct value
        mining_limit = next_block_size_limit
        gbtl = node.getblocktemplatelight({})
        mi = node.getmininginfo()
        assert_greater_than(gbtl['sizelimit'], DEFAULT_CONSENSUS_BLOCK_SIZE)
        assert_equal(mi['miningblocksizelimit'], mining_limit)
        assert_equal(gbtl['sizelimit'], next_block_size_limit)
        # Check sigop limit of the block
        expected_sigops = next_block_size_limit // BLOCK_MAXBYTES_MAXSIGCHECKS_RATIO
        assert_equal(gbtl['sigoplimit'], expected_sigops)

        # Next, restart node, specifying -precentmaxblocksize=50. `getmininginfo` limits should reflect this
        self.restart_node(0, extra_args=self.base_extra_args + [f'-upgrade10activationheight={activation_height}',
                                                                '-percentblockmaxsize=50'])

        mining_limit = next_block_size_limit // 2
        # Check that mempool > 1/2 the blocksize limit so below test works ...
        assert_greater_than_or_equal(node.getmempoolinfo()['bytes'], mining_limit)
        gbtl = node.getblocktemplatelight({})
        mi = node.getmininginfo()
        # GBT doesn't tell you the configured size limit, just the consensus limit :(
        assert_equal(gbtl['sizelimit'], next_block_size_limit)
        # Mining info is the one that tells you the blockmaxsize setting...
        assert_equal(mi['miningblocksizelimit'], mining_limit)
        assert_greater_than_or_equal(mining_limit, mi['currentblocksize'])
        reserved_space = 1000  # BlockAssembler reserves 1000 bytes for coinbase txn
        assert_greater_than_or_equal(mi['currentblocksize'] + reserved_space, mining_limit)
        # Finally, belt-and-suspenders check sigop limit of the block (should be unaffected by blockmaxsize)
        expected_sigops = next_block_size_limit // BLOCK_MAXBYTES_MAXSIGCHECKS_RATIO
        assert_equal(gbtl['sigoplimit'], expected_sigops)

        # Mine beyond the initial 32 MB limit to test ABLA actually allowing for bigger blocks.
        self.restart_node(0, extra_args=self.base_extra_args + [f'-upgrade10activationheight={activation_height}',
                                                                '-percentblockmaxsize=100'])
        assert_greater_than_or_equal(node.getmempoolinfo()['bytes'], DEFAULT_CONSENSUS_BLOCK_SIZE * 2)
        new_block_hash = node.generatetoaddress(2, addr)[1]  # Mine 2 full blocks to bump ABLA state >32MB
        assert_equal(new_block_hash, node.getbestblockhash())  # Ensure blocks actually accepted & connected
        new_block_size = node.getblock(new_block_hash, 1)['size']
        new_abla_state = new_abla_state.NextBlockState(new_block_size)  # Advance abla state
        self.log.info(f"Mined block size: {new_block_size}")
        assert_greater_than(new_abla_state.GetBlockSizeLimit(), DEFAULT_CONSENSUS_BLOCK_SIZE)
        assert_greater_than(new_abla_state.GetNextBlockSizeLimit(), DEFAULT_CONSENSUS_BLOCK_SIZE)
        assert_greater_than(new_block_size, DEFAULT_CONSENSUS_BLOCK_SIZE)

        # "Abusing" the node by restarting it again with upgrade10 activation at the first block height should make the
        # node figure out how to recover and have the correct abla state for the activation block, which is now genesis.
        self.restart_node(0, extra_args=self.base_extra_args + ['-upgrade10activationheight=0'])

        assert_equal(self.get_abla_activation_height(), 0)
        genesis_size = node.getblock(node.getblockhash(0), 1)['size']
        # Genesis abla state should be "default" AblaState (with the blockSize applied)
        abla = AblaState.FromRpcDict(self.get_header(0))
        assert_equal(abla, AblaState(blockSize=genesis_size))

        # Iterate through all blocks and ensure that the entire chain got the corrected abla state post-restart
        last_height = self.get_header()['height']
        for ht in range(1, last_height + 1):
            bh = node.getblockhash(ht)
            size = node.getblock(bh, 1)['size']
            new_abla = abla.NextBlockState(size)
            assert_equal(AblaState.FromRpcDict(self.get_header(bh)), new_abla)
            abla = new_abla

        # Lastly, connect the nodes and ensure node0 and node1 synch up (ABLA working ok to synch blocks)
        connect_nodes(self.nodes[0], self.nodes[1])
        self.sync_blocks(timeout=self.rpc_timeout)


if __name__ == '__main__':
    AblaTest().main()
