from __future__ import annotations

import socket
import threading
import unittest

from bot_bottle.backend.smolmachines.port_forward import (
    ForwardSpec,
    start_forwarder,
    stop_forwarder,
)


def _echo_server() -> tuple[socket.socket, int]:
    listener = socket.create_server(("127.0.0.1", 0))
    port = int(listener.getsockname()[1])

    def run() -> None:
        with listener:
            conn, _ = listener.accept()
            with conn:
                while True:
                    data = conn.recv(65536)
                    if not data:
                        return
                    conn.sendall(data)

    threading.Thread(target=run, daemon=True).start()
    return listener, port


class TestForwarderValidation(unittest.TestCase):
    def test_rejects_generic_localhost_listener(self):
        with self.assertRaises(SystemExit):
            start_forwarder((
                ForwardSpec(
                    label="egress",
                    listen_host="127.0.0.1",
                    listen_port=0,
                    target_host="127.0.0.1",
                    target_port=9099,
                ),
            ))

    def test_rejects_wildcard_listener(self):
        with self.assertRaises(SystemExit):
            start_forwarder((
                ForwardSpec(
                    label="egress",
                    listen_host="0.0.0.0",
                    listen_port=0,
                    target_host="127.0.0.1",
                    target_port=9099,
                ),
            ))

    def test_rejects_non_loopback_target(self):
        with self.assertRaises(SystemExit):
            start_forwarder((
                ForwardSpec(
                    label="egress",
                    listen_host="127.0.0.16",
                    listen_port=0,
                    target_host="192.0.2.10",
                    target_port=9099,
                ),
            ))


class TestForwarding(unittest.TestCase):
    def test_forwards_bytes_both_directions(self):
        try:
            probe = socket.create_server(("127.0.0.16", 0))
        except OSError:
            self.skipTest("127.0.0.16 is not bound on this host")
        else:
            probe.close()
        _listener, target_port = _echo_server()
        handle = start_forwarder((
            ForwardSpec(
                label="egress",
                listen_host="127.0.0.16",
                listen_port=0,
                target_host="127.0.0.1",
                target_port=target_port,
            ),
        ))
        try:
            bound = handle.forwards[0]
            with socket.create_connection((bound.listen_host, bound.listen_port)) as sock:
                sock.sendall(b"hello")
                self.assertEqual(b"hello", sock.recv(5))
        finally:
            stop_forwarder(handle)


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