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

DUST = 546


class ElectrumTokenListUnspentTests(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(110)

        cli = ElectrumConnection()
        await cli.connect()
        await self.sync_height(cli)
        try:
            await self.test_listunspent(n, cli)
            await self.test_listunspent_filtered(n, cli)
            await self.test_error_on_invalid_cursor(n, cli)
            await self.test_listunspent_nft(n, cli)
        finally:
            cli.disconnect()

    async def test_listunspent(self, n, cli):
        addr = n.getnewaddress()
        addr_mint = n.getnewaddress()
        addr_scripthash = await cli.call("blockchain.address.get_scripthash", addr)
        utxo = (await cli.call("token.address.listunspent", addr))["unspent"]
        assert_equal(0, len(utxo))
        assert_equal(
            utxo,
            (await cli.call("token.scripthash.listunspent", addr_scripthash))[
                "unspent"
            ],
        )

        token1_id = self.create_token(to_addr=addr_mint, mint_amount=100)

        txidem = self.send_token(token1_id, addr, 42)
        txid = await get_txid_from_idem(n, txidem)

        async def fetch_utxo():
            print(await cli.call("token.address.listunspent", addr))
            utxo = (await cli.call("token.address.listunspent", addr))["unspent"]
            if len(utxo) > 0:
                return utxo
            return None

        utxo = await wait_for(10, fetch_utxo)
        print(utxo)
        assert_equal(1, len(utxo))

        assert_equal(0, utxo[0]["height"])
        assert_equal(txid, utxo[0]["tx_hash"])
        if on_bch():
            # The 1000 amount is from test_framework.bch.utiltoken
            assert_equal(1000, utxo[0]["value"])
            assert "token_id" in utxo[0]
        elif on_nex():
            assert_equal(DUST, utxo[0]["value"])
            assert "token_id_hex" in utxo[0]
            assert_equal(token1_id, utxo[0]["group"])
        else:
            raise NotImplementedError()
        assert_equal(42, utxo[0]["token_amount"])
        assert utxo[0]["tx_pos"] in [0, 1]

        assert_equal(
            utxo,
            (await cli.call("token.scripthash.listunspent", addr_scripthash))[
                "unspent"
            ],
        )

        n.generate(1)

        async def wait_for_confheight():
            try:
                utxo = (await cli.call("token.address.listunspent", addr))["unspent"]
                return len(utxo) == 1 and utxo[0]["height"] == n.getblockcount()
            except Exception as e:
                if "expected to find token output" in f"{e}":
                    # This is a known race condition bug; where the utxo has been stored in the
                    # database, but not the token info for the utxo
                    return False
                raise e

        await wait_for(10, wait_for_confheight)

    async def test_listunspent_filtered(self, n, cli):
        addr = n.getnewaddress()
        _token1 = self.create_token(to_addr=addr, mint_amount=100)
        token2 = self.create_token(to_addr=addr, mint_amount=100)
        # Minting tokens with this framework uses 2 transactions.
        await self.wait_for_mempool_count(cli, n, count=4)
        utxos = (await cli.call("token.address.listunspent", addr))["unspent"]
        assert_equal(2, len(utxos))
        utxos = (await cli.call("token.address.listunspent", addr, None, token2))[
            "unspent"
        ]
        assert_equal(1, len(utxos))
        if on_bch():
            assert_equal(token2, utxos[0]["token_id"])
        if on_nex():
            assert_equal(token2, utxos[0]["group"])

    async def test_error_on_invalid_cursor(self, n, cli):
        addr = n.getnewaddress()
        token = self.create_token(to_addr=addr, mint_amount=42)
        await self.sync_height(cli)
        await assert_response_error(
            lambda: cli.call("token.address.listunspent", addr, "invalid"),
            ERROR_CODE_INVALID_PARAMS,
            "cursor",
        )
        await assert_response_error(
            lambda: cli.call("token.address.listunspent", addr, "invalid", token),
            ERROR_CODE_INVALID_PARAMS,
            "cursor",
        )

        # cleanup mempool
        n.generate(1)
        await self.wait_for_mempool_count(cli, n, count=0)

    async def test_listunspent_nft(self, n, cli):
        addr = n.getnewaddress()
        token = self.create_token(to_addr=addr, mint_amount=42, bch_can_mint_nft=True)
        commitment_hex = self.create_nft(token, addr, "42")

        async def check():
            if on_bch():
                unspent = await cli.call("token.address.listunspent", addr)
                assert_equal(1, len(unspent["unspent"]))
                assert_equal(commitment_hex, unspent["unspent"][0]["commitment"])

                unspent = await cli.call("token.address.listunspent", addr, None, token)
                assert_equal(1, len(unspent["unspent"]))
                assert_equal(commitment_hex, unspent["unspent"][0]["commitment"])

            if on_nex():
                unspent = await cli.call("token.address.listunspent", addr)
                assert any(
                    utxo["group"] == commitment_hex for utxo in unspent["unspent"]
                )

        if on_bch():
            await self.wait_for_mempool_count(cli, n, count=2)
            await check()

        if on_nex():
            await self.wait_for_mempool_count(cli, n, count=3)
            await check()

        n.generate(1)
        await self.sync_height(cli, n)
        await self.wait_for_mempool_count(cli, n, count=0)
        await check()

        # print(n.getrawtransaction(unspent["unspent"][0]["tx_hash"]))


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