from __future__ import annotations

import socket
import unittest

from straywild_discovery.relay import (
    RELAY_AUTH_PREFIX,
    RelayConfig,
    RelayService,
)


class RelayServiceTests(unittest.TestCase):
    def setUp(self) -> None:
        self.host = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.host.bind(("127.0.0.1", 0))
        self.host.settimeout(0.35)
        self.client = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.client.bind(("127.0.0.1", 0))
        self.client.settimeout(1.0)
        self.relay = RelayService(RelayConfig(
            bind_host="127.0.0.1",
            public_host="127.0.0.1",
            port_start=0,
            port_end=0,
            max_allocations=4,
        ))
        self.relay.start()

    def tearDown(self) -> None:
        self.relay.close()
        self.client.close()
        self.host.close()

    def test_authenticated_allocation_masks_both_peer_endpoints(self) -> None:
        token = "join-capability"
        route = self.relay.allocate(
            token,
            "127.0.0.1",
            int(self.host.getsockname()[1]),
        )
        relay_endpoint = (route.host, route.port)

        self.client.sendto(b"not authenticated", relay_endpoint)
        with self.assertRaises(TimeoutError):
            self.host.recvfrom(4096)

        self.client.sendto(RELAY_AUTH_PREFIX + token.encode(), relay_endpoint)
        self.client.sendto(b"client gameplay", relay_endpoint)
        packet, relay_as_seen_by_host = self.host.recvfrom(4096)
        self.assertEqual(packet, b"client gameplay")
        self.assertEqual(relay_as_seen_by_host[1], route.port)
        self.assertNotEqual(
            relay_as_seen_by_host[1], int(self.client.getsockname()[1])
        )

        self.host.sendto(b"host gameplay", relay_as_seen_by_host)
        packet, relay_as_seen_by_client = self.client.recvfrom(4096)
        self.assertEqual(packet, b"host gameplay")
        self.assertEqual(relay_as_seen_by_client, relay_endpoint)
        self.assertNotEqual(
            relay_as_seen_by_client[1], int(self.host.getsockname()[1])
        )

    def test_wrong_capability_cannot_claim_allocation(self) -> None:
        route = self.relay.allocate(
            "right-token",
            "127.0.0.1",
            int(self.host.getsockname()[1]),
        )
        relay_endpoint = (route.host, route.port)
        self.client.sendto(RELAY_AUTH_PREFIX + b"wrong-token", relay_endpoint)
        self.client.sendto(b"blocked gameplay", relay_endpoint)
        with self.assertRaises(TimeoutError):
            self.host.recvfrom(4096)


if __name__ == "__main__":
    unittest.main()