#!/usr/bin/env python3
# Copyright (c) 2021-2023 The Bitcoin Unlimited developers
import asyncio

from test_framework.utilcompat import lock, unlock
from test_framework.serialize import to_hex
from test_framework.script import CScript, OP_DROP, OP_NOP, OP_TRUE
from test_framework.util import assert_equal
from test_framework.electrumutil import (
    ERROR_CODE_INVALID_PARAMS,
    ElectrumTestFramework,
    assert_response_error,
    script_to_scripthash,
    get_txid_from_idem,
)
from test_framework.environment import on_bch, on_nex, NodeFeature, node_supports, node
from test_framework.electrumconnection import ElectrumConnection
from test_framework.utiltx import pad_tx

if on_bch():
    from test_framework.bch.script import anyonecanspend_unlocking_script
    from test_framework.bch.blocktools import create_transaction
elif on_nex():
    from test_framework.nex.blocktools import create_transaction
    from test_framework.nex.script import anyonecanspend_unlocking_script
else:
    raise NotImplementedError()

GET_UTXO = "blockchain.utxo.get"
TX_BROADCAST = "blockchain.transaction.broadcast"


class ElectrumUtxoTests(ElectrumTestFramework):
    async def run_test(self):

        await self.bootstrap_p2p()
        cli = ElectrumConnection()
        await cli.connect()
        coinbases = await self.mine_blocks(cli, self.nodes[0], 101)

        try:
            await self.test_invalid_txid(cli)
            await self.test_invalid_output(cli, coinbases.pop(0))
            await self.test_two_tx_chain(cli, coinbases.pop(0))
            if node_supports(node(), NodeFeature.TOKENS):
                await self.test_token_output(cli)
        finally:
            cli.disconnect()

    async def test_invalid_txid(self, cli):
        if not on_bch():
            return

        txid = "0000000000000000000000000000000000000000000000000000000000000042"
        await assert_response_error(
            lambda: cli.call(GET_UTXO, txid, 0),
            ERROR_CODE_INVALID_PARAMS,
            "No such mempool transaction",
        )

    async def test_invalid_output(self, cli, unspent):
        if on_nex():
            outpointhash = (
                "0000000000000000000000000000000000000000000000000000000000000042"
            )
            await assert_response_error(
                lambda: cli.call(GET_UTXO, outpointhash),
                ERROR_CODE_INVALID_PARAMS,
                "Outpoint not found",
            )
            return

        if on_bch():
            await assert_response_error(
                lambda: cli.call(GET_UTXO, unspent.rehash(), 420),
                ERROR_CODE_INVALID_PARAMS,
                "out_n 420 does not exist on tx",
            )
            return

        raise NotImplementedError()

    # pylint: disable=too-many-statements
    async def test_two_tx_chain(self, cli, unspent):
        if on_nex():
            res = await cli.call(GET_UTXO, unspent.outpoint_at(0).rpc_hex())
        elif on_bch():
            res = await cli.call(GET_UTXO, unspent.rehash(), 0)
        else:
            raise NotImplementedError()

        assert_equal(res["status"], "unspent")

        original_amount = unspent.vout[0].n_value

        # Use custom locking script for testing scripthash later.
        scriptpubkey1 = CScript([OP_NOP])
        scriptpubkey2 = CScript([OP_TRUE, OP_DROP])

        # Create two transactions, where the second has the first one as
        # parent.
        tx1 = create_transaction(
            unspent,
            n=0,
            sig=anyonecanspend_unlocking_script(),
            out=lock(scriptpubkey1),
            value=original_amount - 1000,
        )
        pad_tx(tx1)
        tx2 = create_transaction(
            tx1,
            n=0,
            sig=unlock(scriptpubkey1),
            out=lock(scriptpubkey2),
            value=original_amount - 2000,
        )
        pad_tx(tx2)

        n = self.nodes[0]
        tx1_id = await get_txid_from_idem(n, await cli.call(TX_BROADCAST, tx1.to_hex()))
        tx2_id = await get_txid_from_idem(n, await cli.call(TX_BROADCAST, tx2.to_hex()))
        await self.wait_for_mempool_count(cli, count=2)

        async def check_utxo(confirmation_height):
            output_index = 0
            if on_nex():
                tx1_utxo = await cli.call(
                    GET_UTXO, tx1.outpoint_at(output_index).rpc_hex()
                )
                tx2_utxo = await cli.call(
                    GET_UTXO, tx2.outpoint_at(output_index).rpc_hex()
                )
            elif on_bch():
                tx1_utxo = await cli.call(GET_UTXO, tx1_id, output_index)
                tx2_utxo = await cli.call(GET_UTXO, tx2_id, output_index)
            else:
                raise NotImplementedError()

            assert_equal(tx1_utxo["status"], "spent")
            assert_equal(tx1_utxo["amount"], original_amount - 1000)
            assert_equal(
                tx1_utxo["scriptpubkey"], to_hex(lock(scriptpubkey1)).decode("ascii")
            )
            assert_equal(
                tx1_utxo["scripthash"], script_to_scripthash(lock(scriptpubkey1))
            )
            assert_equal(tx1_utxo["spent"]["tx_hash"], tx2_id)
            assert_equal(tx1_utxo["spent"]["tx_pos"], 0)
            if confirmation_height is not None:
                assert_equal(tx1_utxo["height"], confirmation_height)
                assert_equal(tx1_utxo["spent"]["height"], confirmation_height)
            else:
                # 0 == in mempool
                assert_equal(tx1_utxo["height"], 0)
                assert_equal(tx1_utxo["spent"]["height"], 0)

            if on_nex():
                assert_equal(tx1_utxo["tx_idem"], tx1.get_rpc_hex_idem())
                assert_equal(tx1_utxo["tx_hash"], tx1.get_rpc_hex_id())
                assert_equal(tx1_utxo["tx_pos"], output_index)

            assert_equal(tx2_utxo["status"], "unspent")
            assert_equal(tx2_utxo["amount"], original_amount - 2000)
            assert_equal(
                tx2_utxo["scriptpubkey"], to_hex(lock(scriptpubkey2)).decode("ascii")
            )
            assert_equal(
                tx2_utxo["scripthash"], script_to_scripthash(lock(scriptpubkey2))
            )
            assert_equal(tx2_utxo["spent"]["height"], None)
            assert_equal(tx2_utxo["spent"]["tx_hash"], None)
            assert_equal(tx2_utxo["spent"]["tx_pos"], None)
            if confirmation_height is not None:
                assert_equal(tx2_utxo["height"], confirmation_height)
            else:
                assert_equal(
                    tx2_utxo["height"], -1
                )  # in mempool, with parent in mempool

            if on_nex():
                assert_equal(tx2_utxo["tx_idem"], tx2.get_rpc_hex_idem())
                assert_equal(tx2_utxo["tx_hash"], tx2.get_rpc_hex_id())
                assert_equal(tx2_utxo["tx_pos"], output_index)

        # Check result when unconfirmed
        await check_utxo(confirmation_height=None)

        # Check result when confirmed
        await self.mine_blocks(cli, self.nodes[0], 1, [tx1, tx2])
        await check_utxo(confirmation_height=self.nodes[0].getblockcount())

    async def create_token_tx(self, n, cli):
        addr = n.getnewaddress()
        _, mint_txid = await self.create_token(to_addr=addr, mint_amount=42)
        await self.wait_for_mempool_tx(cli, mint_txid)
        return (mint_txid, n.getrawtransaction(mint_txid, True), addr)

    async def test_token_output(self, cli):
        # Use the node wallet to create token transactions
        n = self.nodes[0]
        n.generate(101)
        txid, tx, addr = await self.create_token_tx(n, cli)
        if on_bch():
            token_output = list(filter(lambda x: "tokenData" in x, tx["vout"]))[0]
            utxo = await cli.call(GET_UTXO, txid, token_output["n"])
            assert addr in utxo["addresses"]
            assert_equal(token_output["tokenData"]["category"], utxo["token_id"])
            assert_equal(int(token_output["tokenData"]["amount"]), utxo["token_amount"])
            assert "token_bitfield" in utxo
        if on_nex():
            token_output = list(
                filter(lambda x: "group" in x["scriptPubKey"], tx["vout"])
            )[0]
            utxo = await cli.call(GET_UTXO, token_output["outpoint"])
            assert_equal([addr], utxo["addresses"])
            assert_equal(token_output["scriptPubKey"]["group"], utxo["group"])
            assert_equal(
                token_output["scriptPubKey"]["groupQuantity"],
                utxo["group_quantity"],
            )
            assert_equal(
                token_output["scriptPubKey"]["scriptHash"].lower(),
                utxo["template_scripthash"],
            )
            assert_equal(
                token_output["scriptPubKey"]["argsHash"].lower(),
                utxo["template_argumenthash"],
            )
            assert_equal(
                token_output["scriptPubKey"]["groupAuthority"],
                utxo["group_authority"],
            )


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