from __future__ import annotations

import json
import socket
import subprocess
import threading
import unittest
from unittest.mock import MagicMock, patch

from bot_bottle.backend.smolmachines.port_forward import (
    ForwardSpec,
    ForwarderHandle,
    _CompletedProcess,
    _Forwarder,
    _pipe,
    main,
    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_empty_specs_returns_noop_handle(self):
        handle = start_forwarder(())

        self.assertEqual((), handle.forwards)
        self.assertEqual(0, handle.process.poll())
        stop_forwarder(handle)

    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,
                ),
            ))

    def test_rejects_invalid_ports(self):
        with self.assertRaises(SystemExit):
            start_forwarder((
                ForwardSpec("bad-listen", "127.0.0.2", 70000, "127.0.0.1", 9099),
            ))
        with self.assertRaises(SystemExit):
            start_forwarder((
                ForwardSpec("bad-target", "127.0.0.2", 0, "127.0.0.1", 0),
            ))

    def test_completed_process_methods_are_noops(self):
        proc = _CompletedProcess()

        proc.terminate()
        proc.kill()

        self.assertEqual(0, proc.poll())
        self.assertEqual(0, proc.wait(timeout=1))


class _FakeStdout:
    def __init__(self, line: str):
        self.line = line
        self.closed = False

    def readline(self) -> str:
        return self.line

    def close(self) -> None:
        self.closed = True


class TestForwarderProcess(unittest.TestCase):
    def test_start_parses_ready_payload(self):
        payload = {
            "forwards": [{
                "label": "egress",
                "listen_host": "127.0.0.2",
                "listen_port": 61000,
                "target_host": "127.0.0.1",
                "target_port": 55000,
            }],
        }
        proc = MagicMock()
        proc.stdout = _FakeStdout(json.dumps(payload) + "\n")

        with patch(
            "bot_bottle.backend.smolmachines.port_forward.subprocess.Popen",
            return_value=proc,
        ) as popen:
            handle = start_forwarder((
                ForwardSpec("egress", "127.0.0.2", 0, "127.0.0.1", 55000),
            ))

        self.assertIs(proc, handle.process)
        self.assertEqual(61000, handle.forwards[0].listen_port)
        self.assertTrue(proc.stdout.closed)
        argv = popen.call_args.args[0]
        self.assertIn("--spec-json", argv)

    def test_start_dies_when_helper_exits_before_ready(self):
        proc = MagicMock()
        proc.stdout = _FakeStdout("")

        with patch(
            "bot_bottle.backend.smolmachines.port_forward.subprocess.Popen",
            return_value=proc,
        ):
            with self.assertRaises(SystemExit):
                start_forwarder((
                    ForwardSpec("egress", "127.0.0.2", 0, "127.0.0.1", 55000),
                ))

        proc.terminate.assert_called_once()
        proc.wait.assert_called_once_with(timeout=5)

    def test_start_dies_on_invalid_ready_payload(self):
        proc = MagicMock()
        proc.stdout = _FakeStdout('{"wrong": []}\n')

        with patch(
            "bot_bottle.backend.smolmachines.port_forward.subprocess.Popen",
            return_value=proc,
        ):
            with self.assertRaises(SystemExit):
                start_forwarder((
                    ForwardSpec("egress", "127.0.0.2", 0, "127.0.0.1", 55000),
                ))

        proc.terminate.assert_called_once()
        proc.wait.assert_called_once_with(timeout=5)

    def test_stop_kills_process_after_timeout(self):
        proc = MagicMock()
        proc.poll.return_value = None
        proc.wait.side_effect = [subprocess.TimeoutExpired("cmd", 5), 0]

        stop_forwarder(ForwarderHandle(process=proc, forwards=()))

        proc.terminate.assert_called_once()
        proc.kill.assert_called_once()
        self.assertEqual(2, proc.wait.call_count)


class TestForwarding(unittest.TestCase):
    def test_forwards_bytes_both_directions(self):
        try:
            probe = socket.create_server(("127.0.0.2", 0))
        except OSError:
            self.skipTest("127.0.0.2 is not bindable on this host")
        else:
            probe.close()
        _listener, target_port = _echo_server()
        handle = start_forwarder((
            ForwardSpec(
                label="egress",
                listen_host="127.0.0.2",
                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)


class TestForwarderInProcess(unittest.TestCase):
    def test_forwarder_serves_and_stops_in_process(self):
        _listener, target_port = _echo_server()
        forwarder = _Forwarder((
            ForwardSpec("egress", "127.0.0.2", 0, "127.0.0.1", target_port),
        ))
        forwarder.start()
        thread = threading.Thread(target=forwarder.serve, daemon=True)
        thread.start()
        try:
            bound = forwarder.bound[0]
            with socket.create_connection((bound.listen_host, bound.listen_port)) as sock:
                sock.sendall(b"ping")
                self.assertEqual(b"ping", sock.recv(4))
        finally:
            forwarder.stop()
            thread.join(timeout=2)

        self.assertFalse(thread.is_alive())

    def test_forwarder_logs_target_connect_failure(self):
        forwarder = _Forwarder((
            ForwardSpec("egress", "127.0.0.2", 0, "127.0.0.1", 1),
        ))
        forwarder.start()
        thread = threading.Thread(target=forwarder.serve, daemon=True)
        thread.start()
        try:
            bound = forwarder.bound[0]
            with patch("bot_bottle.backend.smolmachines.port_forward.warn") as warn:
                with socket.create_connection((bound.listen_host, bound.listen_port)):
                    pass
                for _ in range(100):
                    if warn.called:
                        break
                    threading.Event().wait(0.01)
            self.assertTrue(warn.called)
        finally:
            forwarder.stop()
            thread.join(timeout=2)

    def test_pipe_stops_when_recv_fails(self):
        src, dst = socket.socketpair()
        src.close()
        try:
            _pipe(src, dst)
        finally:
            dst.close()

    def test_pipe_stops_when_send_fails(self):
        src, src_peer = socket.socketpair()
        _dst_peer, dst = socket.socketpair()
        dst.close()
        try:
            src_peer.sendall(b"hello")
            src_peer.shutdown(socket.SHUT_WR)
            _pipe(src, dst)
        finally:
            src.close()
            src_peer.close()
            _dst_peer.close()


class TestMain(unittest.TestCase):
    def test_main_starts_serves_and_stops_forwarder(self):
        calls: list[str] = []
        handlers = []

        class FakeForwarder:
            def __init__(self, specs):  # type: ignore[no-untyped-def]
                self.specs = specs
                self.bound = specs

            def start(self) -> None:
                calls.append("start")

            def serve(self) -> None:
                calls.append("serve")

            def stop(self) -> None:
                calls.append("stop")

        def record_signal(_signum, handler):  # type: ignore[no-untyped-def]
            handlers.append(handler)

        spec = ForwardSpec("egress", "127.0.0.2", 0, "127.0.0.1", 9099)
        with patch(
            "bot_bottle.backend.smolmachines.port_forward._Forwarder",
            FakeForwarder,
        ), patch(
            "bot_bottle.backend.smolmachines.port_forward.signal.signal",
            side_effect=record_signal,
        ), patch("builtins.print") as print_mock:
            self.assertEqual(
                0,
                main(("--spec-json", json.dumps([{
                    "label": spec.label,
                    "listen_host": spec.listen_host,
                    "listen_port": spec.listen_port,
                    "target_host": spec.target_host,
                    "target_port": spec.target_port,
                }]))),
            )
            handlers[0](15, None)

        self.assertEqual(["start", "serve", "stop", "stop"], calls)
        self.assertIn("forwards", print_mock.call_args.args[0])


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