#!/usr/bin/env python3
# Copyright (c) 2019 The Bitcoin Unlimited developers
"""
Tests to check if basic electrum server integration works
"""
import asyncio
from test_framework.util import (
    assert_equal,
    assert_raises_async,
    wait_for,
)
from test_framework.electrumutil import (
    ElectrumTestFramework,
)
from test_framework.constants import COIN
from test_framework.environment import on_bch, on_nex
from test_framework.serialize import to_hex
from test_framework.connectrum.exc import ElectrumErrorResponse
from test_framework.electrumconnection import ElectrumConnection

if on_bch():
    from test_framework.bch.nodemessages import CTransaction, CTxIn, COutPoint
elif on_nex():
    from test_framework.nex.nodemessages import CTransaction, CTxIn, COutPoint
else:
    raise NotImplementedError()


class ElectrumBasicTests(ElectrumTestFramework):
    async def run_test(self):
        n = self.nodes[0]

        self.info("Checking that blocks are indexed")
        n.generate(200)

        electrum_client = ElectrumConnection()
        try:
            await electrum_client.connect()
            await self.test_mempoolsync(n, electrum_client)
            await self.test_unknown_method(electrum_client)
            await self.test_invalid_args(electrum_client)
            await self.test_address_balance(n, electrum_client)
        finally:
            electrum_client.disconnect()

    async def test_mempoolsync(self, n, cli):
        # throws on timeout, failing the test
        await self.sync_all(cli, n)

        self.info("Check that mempool is communicated")
        n.sendtoaddress(n.getnewaddress(), 10)
        assert_equal(1, len(n.getrawtxpool()))
        await self.wait_for_mempool_count(cli, count=1)

        n.generate(1)
        await self.sync_height(cli)
        assert_equal(0, len(n.getrawtxpool()))
        await self.wait_for_mempool_count(cli, count=0)

    async def test_unknown_method(self, electrum_client):
        await assert_raises_async(
            ElectrumErrorResponse,
            electrum_client.call,
            "method.that.does.not.exist",
            42,
        )

        try:
            await electrum_client.call("method.that.does.not.exist")
        except ElectrumErrorResponse as e:
            error_code = "-32601"
            assert error_code in str(e)

    async def test_invalid_args(self, electrum_client):
        error_code = "-32602"

        hash_param_methods = (
            "blockchain.scripthash.get_balance",
            "blockchain.scripthash.get_history",
            "blockchain.scripthash.listunspent",
        )

        for method in hash_param_methods:
            await assert_raises_async(
                ElectrumErrorResponse, electrum_client.call, method, "invalidhash"
            )
            try:
                await electrum_client.call(method, "invalidhash")
            except ElectrumErrorResponse as e:
                print("ERROR:" + str(e))
                assert error_code in str(e)

        # invalid tx
        try:
            tx = CTransaction()
            if on_nex():
                tx.vin = [CTxIn(COutPoint().from_idem_and_idx(0xF00D, 0))]
            else:
                tx.vin = [CTxIn(COutPoint(0xF00D, 0))]
            tx.rehash()
            await electrum_client.call("blockchain.transaction.broadcast", to_hex(tx))
        except Exception as e:
            assert error_code in str(e), e

    async def test_address_balance(self, n, electrum_client):
        addr = n.getnewaddress()
        n.sendtoaddress(addr, 10)

        async def check_address(address, unconfirmed=0, confirmed=0):
            res = await electrum_client.call("blockchain.address.get_balance", address)
            print(res, "vs", unconfirmed * COIN)

            return (
                res["unconfirmed"] == unconfirmed * COIN
                and res["confirmed"] == confirmed * COIN
            )

        await wait_for(10, lambda: check_address(addr, unconfirmed=10))
        n.generate(1)
        await wait_for(10, lambda: check_address(addr, confirmed=10))


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