#!/usr/bin/env python3
# Copyright (c) 2024 The Bitcoin Unlimited developers

"""
Tests for block height filtering
"""

import asyncio
from test_framework.electrumutil import (
    ElectrumTestFramework,
    get_txid_from_idem,
)
from test_framework.electrumconnection import ElectrumConnection


def is_tx_hash_in_history(history, tx_hash):
    return any(entry.get("tx_hash") == tx_hash for entry in history)


def get_history_height(history, tx_hash):
    for entry in history:
        if entry.get("tx_hash") == tx_hash:
            return entry.get("height")
    assert False, "entry not found"


# pylint: disable=too-many-locals, too-many-arguments
class ElectrumTokenHistoryTests(ElectrumTestFramework):
    async def run_test(self):
        # This test users nexad wallet to create and send tokens.
        # Mine and mature some coins.
        n = self.nodes[0]
        n.generate(101)

        cli = ElectrumConnection()
        await cli.connect()
        try:
            await self.sync_height(cli)

            gen_token_addr = n.getnewaddress()
            token_id, mint_txid = await self.create_token(
                to_addr=gen_token_addr, mint_amount=1000
            )
            print("mint txid", mint_txid)
            n.generate(1)
            await self.sync_all(cli)

            # Test setup; 2 transactions to address at first_height (1 token, 1 regular)
            addr = n.getnewaddress()
            first_height = n.getblockcount() + 1
            txid1 = await get_txid_from_idem(n, n.sendtoaddress(addr, 21))
            txid2 = await self.send_token(token_id, addr, 21)
            await self.wait_for_mempool_count(cli, count=2)
            n.generate(10)
            await self.sync_all(cli)

            block1 = {
                "height": first_height,
                "regular": txid1,
                "token": txid2,
            }

            # 2 transaction at next block
            second_height = first_height + 10
            txid3 = await get_txid_from_idem(n, n.sendtoaddress(addr, 21))
            txid4 = await self.send_token(token_id, addr, 21)
            await self.wait_for_mempool_count(cli, count=2)
            n.generate(1)
            await self.sync_all(cli)

            block2 = {"height": second_height, "regular": txid3, "token": txid4}

            # 2 transactions in mempool
            txid5 = await get_txid_from_idem(n, n.sendtoaddress(addr, 21))
            txid6 = await self.send_token(token_id, addr, 21)
            await self.wait_for_mempool_count(cli, count=2)

            mempool = {"height": 0, "regular": txid5, "token": txid6}

            print("FIRST HEIGHT:", first_height)
            print("SECOND HEIGHT:", second_height)

            async def run_height_filter_tests(test_method):
                print("Testing {}", test_method)
                await test_method(
                    cli,
                    addr,
                    [block1, block2, mempool],
                    from_height=None,
                    to_height=None,
                )
                await test_method(
                    cli, addr, [block1], from_height=0, to_height=first_height + 1
                )
                await test_method(
                    cli,
                    addr,
                    [block1, block2],
                    from_height=0,
                    to_height=second_height + 1,
                )

                await test_method(
                    cli,
                    addr,
                    [block1],
                    from_height=first_height,
                    to_height=first_height + 1,
                )
                await test_method(
                    cli,
                    addr,
                    [block1, block2],
                    from_height=first_height,
                    to_height=second_height + 1,
                )
                await test_method(
                    cli,
                    addr,
                    [block2],
                    from_height=second_height,
                    to_height=second_height + 1,
                )

                # include mempool
                await test_method(
                    cli,
                    addr,
                    [block1, block2, mempool],
                    from_height=first_height,
                    to_height=-1,
                )
                await test_method(
                    cli,
                    addr,
                    [block2, mempool],
                    from_height=second_height,
                    to_height=-1,
                )
                await test_method(cli, addr, [mempool], from_height=1000, to_height=-1)

            await run_height_filter_tests(self.test_history_height_filter)
            await run_height_filter_tests(self.test_token_history_height_filter)

        finally:
            cli.disconnect()

    async def test_history_height_filter(
        self, cli, addr, expected, *, from_height: int | None, to_height: int | None
    ):

        query_filter = {"from_height": from_height, "to_height": to_height}

        history = await cli.call("blockchain.address.get_history", addr, query_filter)
        history_len = len(history)

        for block in expected:
            assert is_tx_hash_in_history(
                history, block["regular"]
            ), f"Did not find regular {block['regular']} at height {block['height']}"
            assert is_tx_hash_in_history(
                history, block["token"]
            ), f"Did not find token {block['token']} at height {block['height']}"

            expected_blocks = [block["height"]]
            if block["height"] == 0:
                # allow it to be -1 also (unconfirmed parent)
                expected_blocks.append(-1)

            assert (
                get_history_height(history, block["regular"]) in expected_blocks
            ), f"height wrong for {block["regular"]}, expected {block["height"]}, got {get_history_height(history, block["regular"])}"
            assert (
                get_history_height(history, block["token"]) in expected_blocks
            ), f"height wrong for {block["token"]}, expected {block["height"]}, got {get_history_height(history, block["token"])}"

        expected_len = len(expected) * 2  # regular + token
        assert (
            history_len == expected_len
        ), f"Expected {expected_len} history entries, got {history_len}"

        print("OK")

    async def test_token_history_height_filter(
        self, cli, addr, expected, *, from_height: int | None, to_height: int | None
    ):

        query_filter = {
            "from_height": from_height,
            "to_height": to_height,
            "tokens_only": True,
        }

        history = (
            await cli.call("token.address.get_history", addr, None, None, query_filter)
        )["transactions"]
        history_len = len(history)

        for block in expected:
            assert is_tx_hash_in_history(
                history, block["token"]
            ), f"Did not find token {block['token']} at height {block['height']}"

            expected_blocks = [block["height"]]
            if block["height"] == 0:
                # allow it to be -1 also (unconfirmed parent)
                expected_blocks.append(-1)

            assert (
                get_history_height(history, block["token"]) in expected_blocks
            ), f"height wrong for {block["token"]}, expected {block["height"]}, got {get_history_height(history, block["token"])}"

        expected_len = len(expected)
        assert (
            history_len == expected_len
        ), f"Expected {expected_len} history entries, got {history_len}"

        print("OK")


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