#!/usr/bin/env python3
# Copyright (c) 2019 The Bitcoin Unlimited developers
"""
Tests to check if basic electrum server integration works
"""
import asyncio
from test_framework.util import assert_equal
from test_framework.electrumutil import (
    ElectrumTestFramework,
)
from test_framework.environment import node, node_supports, NodeFeature
from test_framework.electrumconnection import ElectrumConnection

DEFAULT_WAIT = 30


def confirmed_count(history):
    # token.* adds the history list under 'history' key,
    # while address.* responses do not.
    if "history" in history:
        history = history["history"]

    return sum(map(lambda i: i["height"] > 0, history))


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

        n.generate(200)
        cli = ElectrumConnection()
        await cli.connect()
        try:
            await self.test_reorg(cli, n)
        finally:
            cli.disconnect()

    async def test_reorg(self, cli, n):
        address = n.getnewaddress()

        await self.sync_height(cli)
        await self.wait_for_mempool_count(cli, count=0)

        test_tokens = node_supports(node(), NodeFeature.TOKENS)

        # 1 coin transaction
        n.sendtoaddress(address, 10)
        await self.wait_for_mempool_count(cli, count=1)

        if test_tokens:
            # 2 token transactions
            token_id = self.create_token(to_addr=n.getnewaddress(), mint_amount=42)
            await self.wait_for_mempool_count(cli, count=3)
        await self.sync_mempool_count(cli)

        blocks = n.generate(50)

        await self.sync_height(cli)
        await self.wait_for_mempool_count(cli, count=0)

        # Check that our transactions got indexed.
        assert (
            confirmed_count(await cli.call("blockchain.address.get_history", address))
            >= 1
        )
        if test_tokens:
            assert_equal(
                2,
                confirmed_count(
                    await cli.call("token.transaction.get_history", token_id)
                ),
            )

        self.info("invalidating %d blocks", len(blocks))
        n.invalidateblock(blocks[0])

        # electrum server should trim its chain as well and maybe see our
        # transactions go back into mempool (depending on maturity of inputs)
        await self.sync_all(cli)

        # the transactions should no longer exist on-chain
        assert_equal(
            0,
            confirmed_count(await cli.call("blockchain.address.get_history", address)),
        )
        if test_tokens:
            assert_equal(
                0,
                confirmed_count(
                    await cli.call("token.transaction.get_history", token_id)
                ),
            )

        n.generate(50)
        await self.sync_height(cli)


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