#!/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
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_height(cli)
        try:
            await self.test_balance(n, cli)
            if on_nex():
                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 = self.create_token(to_addr=addr, mint_amount=100)
        token1_id_hex = await self.decode_groupid(cli, token1_id)

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

        await self.sync_mempool_count(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()
        self.send_token(token1_id, addr3, 100)
        self.send_token(token2_id, addr3, 200)
        self.send_token(token1_id, addr4, 100)
        self.send_token(token2_id, addr4, 200)
        await self.sync_mempool_count(cli)
        b = await cli.call("token.address.get_balance", addr3)
        assert_equal(0, b["unconfirmed"][token1_id_hex])
        assert_equal(0, b["unconfirmed"][token2_id_hex])

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

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

    # Authority test
    async def test_auth_balance(self, n, cli):
        # Check authorities not included in balance
        addr1 = n.getnewaddress()
        addr2 = n.getnewaddress()
        nex_token_id = self.create_token(addr1, to_addr=addr2, mint_amount=100)
        nex_token_id_hex = await self.decode_groupid(cli, nex_token_id)

        await self.sync_mempool_count(cli)
        b = await cli.call("token.address.get_balance", addr1)
        assert_equal(0, b["unconfirmed"][nex_token_id_hex])

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


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