#!/usr/bin/env python3
# Copyright (c) 2025 The Bitcoin Unlimited developers
"""
Tests for electrum peer discovery (servers discovering each other)
"""

import asyncio
import os

from test_framework.environment import (
    testing_tcp_ssl,
    testing_websocket_ssl,
    testing_https,
)
from test_framework.dummyelectrumserver import Protocol, dummy_electrum
from test_framework.util import assert_equal, wait_for
from test_framework.electrumutil import (
    ElectrumTestFramework,
)

os.environ["ROSTRUM_ANNOUNCE"] = "1"


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

        # workaround; rostrum crashes on restart with empty chain
        self.nodes[0].generate(1)
        await self.test_node_accepts()
        await self.test_switch_to_hostname()
        await self.test_invalid_ports()
        await self.test_tcp_and_ssl()
        await self.remote_announces_to_us()

    async def announce_ourself(self, cli, dummy, features):
        # Announce the server to the real electrum server
        resp = await cli.call("server.add_peer", features)
        assert resp, "remote server rejected us"
        self.info("Remote peer accepted our add request")

        self.info("waiting for server to connect to us...")
        await wait_for(60, lambda: dummy.got_remote_connection)
        self.info("remote peer connected")

        # Check that remote server accepted us
        async def has_accepted_us():
            res = await cli.call("server.peers.subscribe")
            return len(res) > 0

        await wait_for(10, has_accepted_us)
        self.info("remote server added us")

    async def test_node_accepts(self):
        async with dummy_electrum() as (dummy, cli):

            await self.announce_ourself(cli, dummy, dummy.features())

            # Verify that our info is correct
            res = await cli.call("server.peers.subscribe")
            ip_addr, hostname, meta = res.pop()
            assert_equal("127.0.0.1", ip_addr)
            assert_equal("localhost", hostname)
            assert f"v{dummy.protocol_max}" in meta
            assert f"t{dummy.tcp_port}" in meta

    async def test_switch_to_hostname(self):
        """
        Test that even though candidate was initially added without hostname; that
        after asking the peer for features, we use that hostname
        """

        await self.restart_node(0)

        async with dummy_electrum() as (dummy, cli):
            features = dummy.features()
            # replacing with ip instead of 'localhost'
            features["hosts"] = {"127.0.0.1": {"tcp_port": dummy.tcp_port}}
            await self.announce_ourself(cli, dummy, features)

            # Verify that our info is correct
            res = await cli.call("server.peers.subscribe")
            ip_addr, hostname, meta = res.pop()
            assert_equal("127.0.0.1", ip_addr)
            assert_equal("localhost", hostname)
            assert f"v{dummy.protocol_max}" in meta, meta
            assert f"t{dummy.tcp_port}" in meta, meta

    async def test_invalid_ports(self):
        """
        Test that remote server still accepts us; even though we announce invalid ssl or tcp port.
        It should not forward the invalid port to other peers.
        """

        # Claiming to support SSL; but only supporting TCP
        await self.restart_node(0)
        async with dummy_electrum(protocols=[Protocol.TCP]) as (dummy, cli):
            features = dummy.features()
            features["hosts"] = {
                "127.0.0.1": {"ssl_port": 12345, "tcp_port": dummy.tcp_port}
            }
            await self.announce_ourself(cli, dummy, features)

            res = await cli.call("server.peers.subscribe")
            _, _, meta = res.pop()
            assert f"t{dummy.tcp_port}" in meta
            assert not "s12345" in meta

        # Claiming to support TCP; but only supporting SSL
        await self.restart_node(0)
        async with dummy_electrum(protocols=[Protocol.SSL]) as (dummy, cli):
            features = dummy.features()
            features["hosts"] = {
                "127.0.0.1": {"ssl_port": dummy.ssl_port, "tcp_port": 12345}
            }
            await self.announce_ourself(cli, dummy, features)

            res = await cli.call("server.peers.subscribe")
            _, _, meta = res.pop()
            assert f"s{dummy.ssl_port}" in meta
            assert not "t12345" in meta

    async def test_tcp_and_ssl(self):
        # Supporting both TCP and SSL; announce both
        await self.restart_node(0)
        async with dummy_electrum(protocols=[Protocol.TCP, Protocol.SSL]) as (
            dummy,
            cli,
        ):
            features = dummy.features()
            features["hosts"] = {
                "127.0.0.1": {"ssl_port": dummy.ssl_port, "tcp_port": dummy.tcp_port}
            }
            await self.announce_ourself(cli, dummy, features)

            res = await cli.call("server.peers.subscribe")
            _, _, meta = res.pop()
            assert f"s{dummy.ssl_port}" in meta
            assert f"t{dummy.tcp_port}" in meta

        # If we don't initially announce the ssl port to the, it should still discover it
        await self.restart_node(0)
        async with dummy_electrum(protocols=[Protocol.TCP, Protocol.SSL]) as (
            dummy,
            cli,
        ):
            features = dummy.features()
            # TCP only
            features["hosts"] = {"127.0.0.1": {"tcp_port": dummy.tcp_port}}
            await self.announce_ourself(cli, dummy, features)

            res = await cli.call("server.peers.subscribe")
            _, _, meta = res.pop()
            assert f"s{dummy.ssl_port}" in meta
            assert f"t{dummy.tcp_port}" in meta

        # Seeded with wrong TCP port; should still discover correct through the SSL connection
        await self.restart_node(0)
        async with dummy_electrum(protocols=[Protocol.TCP, Protocol.SSL]) as (
            dummy,
            cli,
        ):
            features = dummy.features()
            # TCP only
            features["hosts"] = {
                "127.0.0.1": {"ssl_port": dummy.ssl_port, "tcp_port": 12345}
            }
            await self.announce_ourself(cli, dummy, features)

            res = await cli.call("server.peers.subscribe")
            _, _, meta = res.pop()
            assert f"s{dummy.ssl_port}" in meta
            assert f"t{dummy.tcp_port}" in meta

    async def remote_announces_to_us(self):
        """
        Check that after we've announced dummy server to rostrum,
        that rostrum announces itself back to us.
        """
        await self.restart_node(0)
        async with dummy_electrum() as (dummy, cli):
            await self.announce_ourself(cli, dummy, dummy.features())
            await wait_for(60, lambda: len(dummy.peers) > 0)
            assert_equal(1, len(dummy.peers))
            announcement = dummy.peers.pop()

            if testing_tcp_ssl() or testing_websocket_ssl() or testing_https():
                # Read from SSL certificate; which in our setup is example.org
                assert "example.org" in announcement["hosts"], announcement
            else:
                assert "localhost" in announcement["hosts"], announcement

            assert "genesis_hash" in announcement
            assert "hash_function" in announcement
            assert "protocol_max" in announcement
            assert "hosts" in announcement
            assert "server_version" in announcement

            # remote should also ask if we know of other peers
            await wait_for(30, lambda: dummy.requested_peers)


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