#!/usr/bin/env python3
# Copyright (c) 2022 The Bitcoin Unlimited developers
import asyncio
from test_framework.environment import on_nex
from test_framework.util import assert_equal, wait_for
from test_framework.electrumutil import (
    ElectrumTestFramework,
)
from test_framework.electrumconnection import ElectrumConnection


class ElectrumTokenGetBalanceTests(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(120)

        cli = ElectrumConnection()
        await cli.connect()
        await self.sync_mempool(cli)
        try:
            await self.test_balance(n, cli)
            await self.test_auth_balance(n, cli)
        finally:
            cli.disconnect()

    # Basic test
    async def test_balance(self, n, cli):
        addr = n.getnewaddress()
        addr_scripthash = await cli.call("blockchain.address.get_scripthash", addr)
        addr2 = n.getnewaddress()

        token1_id, _ = await self.create_token(to_addr=addr, mint_amount=100)
        token1_id_hex = await self.decode_groupid(cli, token1_id)

        token2_id, _ = await self.create_token(to_addr=addr, mint_amount=200)
        token2_id_hex = await self.decode_groupid(cli, token2_id)

        await self.sync_mempool(cli)
        b = await cli.call("token.address.get_balance", addr)
        assert len(b["confirmed"]) == 0
        assert_equal(100, b["unconfirmed"][token1_id_hex])
        assert_equal(200, b["unconfirmed"][token2_id_hex])
        assert_equal(b, await cli.call("token.scripthash.get_balance", addr_scripthash))

        n.generate(1)
        await self.sync_all(cli)
        b = await cli.call("token.address.get_balance", addr)
        assert len(b["unconfirmed"]) == 0
        assert_equal(100, b["confirmed"][token1_id_hex])
        assert_equal(200, b["confirmed"][token2_id_hex])

        # Check filter by token id
        b = await cli.call("token.address.get_balance", addr, None, token1_id)
        assert len(b["confirmed"]) == 1
        assert_equal(100, b["confirmed"][token1_id_hex])

        # Check filter with different token on same address and by token id hex
        b = await cli.call("token.address.get_balance", addr, None, token2_id_hex)
        assert len(b["confirmed"]) == 1
        assert_equal(200, b["confirmed"][token2_id_hex])

        # Special case. Sending and receiving the same
        # amount in either mempool, or confirmed index, results in balance of 0.

        # mempool
        addr3 = n.getnewaddress()
        addr4 = n.getnewaddress()
        await self.send_token(token1_id, addr3, 100)
        await self.send_token(token2_id, addr3, 200)
        await self.send_token(token1_id, addr4, 100)
        await self.send_token(token2_id, addr4, 200)
        await self.sync_mempool(cli)
        b = await cli.call("token.address.get_balance", addr3)
        print("addr3 balance", b)
        print("addr3 history", await cli.call("token.address.get_history", addr3))
        await wait_for(10, lambda: b["unconfirmed"].get(token1_id_hex, 0) == 0)
        await wait_for(10, lambda: b["unconfirmed"].get(token2_id_hex, 0) == 0)

        # confirmed index
        n.generate(1)
        await self.sync_all(cli)
        b = await cli.call("token.address.get_balance", addr3)
        assert_equal(0, b["confirmed"].get(token1_id_hex, 0))
        assert_equal(0, b["confirmed"].get(token2_id_hex, 0))

        # Spending confirmed tokens results in negative mempool
        # balance.
        n.generate(1)
        await self.sync_all(cli)
        await self.send_token(token1_id, addr2, 100)
        await self.send_token(token2_id, addr2, 200)
        await self.sync_mempool(cli)
        b = await cli.call("token.address.get_balance", addr4)
        await wait_for(5, lambda: b["unconfirmed"][token1_id_hex] == -100)
        await wait_for(5, lambda: b["unconfirmed"][token2_id_hex] == -200)

    # Authority test
    async def test_auth_balance(self, n, cli):

        if not on_nex():
            # This is a NEXA specific test
            return

        auth_addr = n.getnewaddress()
        token = n.token("new")
        group_id = token["groupIdentifier"]
        group_hex = await self.decode_groupid(cli, group_id)
        n.token("authority", "create", group_id, auth_addr, "mint")
        await self.sync_mempool(cli)
        unspent = await cli.call("token.address.listunspent", auth_addr)
        assert_equal(1, len(unspent["unspent"]))

        # The "token amount" in the utxo is negative, representing the authority flags
        assert unspent["unspent"][0]["token_amount"] < 0
        # The balance should be zero still
        balance = await cli.call("token.address.get_balance", auth_addr)
        assert_equal(0, balance["unconfirmed"][group_hex])

        # Check that balance is correct when authority address also contains tokens
        n.token("mint", group_id, auth_addr, 100)
        await self.sync_mempool(cli)
        balance = await cli.call("token.address.get_balance", auth_addr)
        assert_equal(100, balance["unconfirmed"][group_hex])


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