"""Run a couple of UCSPI client and server tests."""

import abc
import dataclasses
import pathlib
import shlex
import socket
import subprocess
import sys

from typing import Any, Callable


MSG_RESP_HELLO = "a01 hello\n"
MSG_RESP_BYE = "a02 bye\n"


@dataclasses.dataclass(frozen=True)
class Config:
    """Runtime configuration for the UCSPI test runner."""

    bindir: pathlib.Path
    proto: str
    utf8_env: dict[str, str]


class RunnerError(Exception):
    """An error that occurred while preparing for or running the tests."""


class Runner(metaclass=abc.ABCMeta):
    """A helper class for running tests for a single UCSPI protocol."""

    _cfg: Config
    _proto: str

    def __init__(self, cfg: Config, proto: str) -> None:
        """Store the configuration object."""
        self._cfg = cfg
        self._proto = proto

    @property
    def cfg(self) -> Config:
        """Get the configuration for this runner."""
        return self._cfg

    @property
    def proto(self) -> str:
        """Get the name of the UCSPI protocol to test."""
        return self._proto

    @abc.abstractmethod
    def find_listening_address(self) -> list[str]:
        """Find an available protocol-specific address to listen on."""
        raise NotImplementedError(f"{type(self).__name__}.find_listening_address() is abstract")

    @abc.abstractmethod
    def get_listening_socket(self, addr: list[str]) -> socket.socket:
        """Start listening on the specified address."""
        raise NotImplementedError(f"{type(self).__name__}.get_listening_socket() is abstract")

    @abc.abstractmethod
    def get_connected_socket(self, addr: list[str]) -> socket.socket:
        """Connect to the specified address."""
        raise NotImplementedError(f"{type(self).__name__}.get_connected_socket() is abstract")

    @abc.abstractmethod
    def format_local_addr(self, addr: list[str]) -> str:
        """Format an address returned by accept(), etc."""
        raise NotImplementedError(f"{type(self).__name__}.format_local_addr() is abstract")

    @abc.abstractmethod
    def format_remote_addr(self, addr: Any) -> str:
        """Format an address returned by accept(), etc."""
        raise NotImplementedError(f"{type(self).__name__}.format_remote_addr() is abstract")

    def __repr__(self) -> str:
        """Provide a Python-esque representation of the helper."""
        return f"{type(self).__name__}(cfg={self._cfg!r}, proto={self._proto!r})"


HANDLERS: dict[str, type[Runner]] = {}


def with_listener(
    runner: Runner, addr: list[str], tag: str, callback: Callable[[socket.socket], None]
) -> None:
    """Get a listener socket from the runner, invoke the callback."""
    proto = runner.proto
    saddr = runner.format_local_addr(addr)
    print(f"- {tag}: setting up a listening socket at {saddr}")
    try:
        listener = runner.get_listening_socket(addr)
    except RunnerError as err:
        sys.exit(f"{tag}: could not get a {proto} listening socket: {err}")

    try:
        callback(listener)
    finally:
        try:
            listener.close()
        except OSError as err:
            print(f"{tag}: could not close the listening socket at {saddr}: {err}", file=sys.stderr)


def with_connect(
    runner: Runner, addr: list[str], tag: str, callback: Callable[[socket.socket], None]
) -> None:
    """Ask the runner to connect to the specified address, invoke the callback."""
    proto = runner.proto
    saddr = runner.format_local_addr(addr)
    print(f"- {tag}: setting up a socket connected to {saddr}")
    try:
        conn = runner.get_connected_socket(addr)
    except RunnerError as err:
        sys.exit(f"{tag}: could not get a {proto} connected socket: {err}")

    try:
        callback(conn)
    finally:
        try:
            conn.close()
        except OSError as err:
            print(f"{tag}: could not close the socket connected to {saddr}: {err}", file=sys.stderr)


def with_child(
    runner: Runner,
    tag: str,
    cmd: list[str],
    callback: Callable[[subprocess.Popen[str]], None],
    stderr: int | None = None,
) -> None:
    """Spawn a child process, invoke the callback."""
    print(f"- {tag}: starting {shlex.join(cmd)}")
    try:
        with subprocess.Popen(
            cmd,
            bufsize=0,
            encoding="UTF-8",
            env=runner.cfg.utf8_env,
            stdout=subprocess.PIPE,
            stderr=stderr,
        ) as child:
            try:
                callback(child)
            finally:
                if child.poll() is None:
                    ptag = f"the {child.pid} ({cmd[0]}) process"
                    print(f"- {tag}: {ptag} is still active, killing it")
                    try:
                        child.kill()
                    except OSError as err:
                        print(
                            f"{tag}: could not send a kill signal to {ptag}: {err}", file=sys.stderr
                        )
                    print(f"- {tag}: waiting for {ptag} to go away")
                    try:
                        child.wait()
                    except OSError as err:
                        print(f"{tag}: could not wait for {ptag} to end: {err}", file=sys.stderr)
    except (OSError, subprocess.CalledProcessError) as err:
        sys.exit(f"Could not spawn {shlex.join(cmd)}: {err}")


def test_local_spew(runner: Runner, addr: list[str], tag: str, cmd: list[str]) -> None:
    """Test a client program against our listening socket, unidirectional transfer."""

    def run(listener: socket.socket, catproc: subprocess.Popen[str]) -> None:
        """Run the spew test itself."""
        saddr = runner.format_local_addr(addr)
        print(f"- started the client as pid {catproc.pid}")
        print(f"- waiting for a connection at {saddr}")
        try:
            conn, rem_addr = listener.accept()
        except OSError as err:
            sys.exit(f"Could not accept an incoming connection at {saddr}: {err}")
        print(
            f"- accepted a connection at fd {conn.fileno()} from "
            f"{runner.format_remote_addr(rem_addr)}"
        )

        msg = MSG_RESP_HELLO + MSG_RESP_BYE
        try:
            conn.sendall(msg.encode("UTF-8"))
        except OSError as err:
            sys.exit(f"Could not send a message on the {conn} connection at {saddr}: {err}")

        print("- closing the incoming connection")
        try:
            conn.close()
        except OSError as err:
            sys.exit(f"Could not close the {conn} connection at {saddr}: {err}")

        try:
            output, _ = catproc.communicate()
        except OSError as err:
            sys.exit(f"Could not read the output of the {cmd[0]} process: {err}")
        assert isinstance(output, str), repr(output)

        try:
            res = catproc.wait()
        except OSError as err:
            sys.exit(f"Could not wait for the {cmd[0]} process to end: {err}")
        print(f"- client exit code: {res}; output: {output!r}")
        if res != 0:
            sys.exit(f"The {cmd[0]} program exited with non-zero code {res}")

        if output != msg:
            sys.exit(f"Expected {msg!r} as {cmd[0]} output, got {catproc.stdout!r}")

    with_listener(
        runner,
        addr,
        tag,
        lambda listener: with_child(runner, tag, cmd, lambda catproc: run(listener, catproc)),
    )
    print(f"- {cmd[0]} seems fine")


def test_local_cat(runner: Runner, addr: list[str]) -> None:
    """Test the {proto}cat program."""
    proto = runner.proto
    print(f"\n=== Testing {proto}cat")
    if runner.cfg.bindir != pathlib.Path("/usr/bin"):
        print(
            f"- test skipped, {proto}cat will probably not find {runner.cfg.bindir}/{proto}client"
        )
        return

    catpath = runner.cfg.bindir / f"{proto}cat"
    print(f"- will spawn {proto}cat at {catpath} in a while")
    test_local_spew(runner, addr, "test_local_cat", [str(catpath)] + addr)


def test_local_client_spew(runner: Runner, addr: list[str]) -> None:
    """Test {proto}client against our own listening socket."""
    proto = runner.proto
    print(f"\n=== Testing {proto}client against our own listening socket")

    clipath = runner.cfg.bindir / f"{proto}client"
    print(f"- will spawn {proto}client at {clipath} in a while")
    test_local_spew(
        runner,
        addr,
        "test_local_client_spew",
        [str(clipath), "-R", "--"] + addr + ["sh", "-c", "set -e; exec <&6; exec cat"],
    )


def test_server_local(runner: Runner, addr: list[str]) -> None:
    """Test {proto}server against our own client socket."""

    def run(srvproc: subprocess.Popen[str], conn: socket.socket) -> None:
        """Run the test itself."""
        try:
            rem_addr = conn.getpeername()
        except OSError as err:
            sys.exit(f"getpeername() failed for {conn!r}: {err}")
        raddr = runner.format_remote_addr(rem_addr)
        print(f"- got a connection at fd {conn.fileno()} to {raddr}")
        print("- reading all the data we can")
        try:
            data = conn.recv(4096)
        except OSError as err:
            sys.exit(f"Could not read from the {conn!r} socket: {err}")
        print(f"- read {len(data)} bytes from the socket")
        try:
            output = data.decode("UTF-8")
        except ValueError as err:
            sys.exit(f"Could not decode {data!r} as valid UTF-8: {err}")
        if output != message:
            sys.exit(f"Expected {message!r}, got {output!r}")

        if srvproc.poll() is not None:
            sys.exit(
                f"Did not expect the {proto}server process to have exited: "
                f"code {srvproc.returncode}"
            )
        print(f"- sending a SIGTERM signal to the {proto}server process")
        try:
            srvproc.terminate()
        except OSError as err:
            sys.exit(
                f"Could not send a SIGTERM signal to the {proto}server process at "
                f"{srvproc.pid}: {err}"
            )
        print(f"- waiting for the {proto}server process to exit")
        res = srvproc.wait()
        if res != 0:
            sys.exit(f"The {proto}server process exited with a non-zero code {res}")

    def wait_and_connect(srvproc: subprocess.Popen[str]) -> None:
        """Wait for the "ready" message from tcpserver."""
        assert srvproc.stderr is not None, repr(srvproc)
        print(f"- awaiting the first 'status' line from {proto}server")
        line = srvproc.stderr.readline()
        if "server: status: 0/" not in line:
            sys.exit(
                f"Unexpected first line from {proto}server: expected 'status: 0/N', got {line!r}"
            )
        with_connect(runner, addr, "test_server_local", lambda conn: run(srvproc, conn))

    proto = runner.proto
    print(f"\n=== Testing {proto}server against our own listening socket")

    srvpath = runner.cfg.bindir / f"{proto}server"
    print(f"- will spawn {proto}client at {srvpath} in a while")
    message = MSG_RESP_HELLO + MSG_RESP_BYE
    with_child(
        runner,
        "test_server_local",
        [str(srvpath), "-vR", "--"] + addr + ["printf", "--", message.replace("\n", "\\n")],
        wait_and_connect,
        stderr=subprocess.PIPE,
    )

    print(f"- test_server_local: {srvpath} seems fine")


def test_server_client_spew(runner: Runner, addr: list[str]) -> None:
    """Test {proto}server against {proto}client, unidirectional data transfer."""

    def run(srvproc: subprocess.Popen[str], cliproc: subprocess.Popen[str]) -> None:
        """Read the data received by the client."""
        try:
            output, _ = cliproc.communicate()
        except OSError as err:
            sys.exit(f"Could not read the output of the {proto}client process: {err}")
        res = cliproc.poll()
        print(f"- client exit code {res}; output {output!r}")
        if res is None:
            sys.exit(f"Expected the {proto}client process to be done by now")
        if res != 0:
            sys.exit(f"The {proto}client process exited with a non-zero code {res}")
        if output != message:
            sys.exit(f"Expected {message!r}, got {output!r}")

        if srvproc.poll() is not None:
            sys.exit(
                f"Did not expect the {proto}server process to have exited: "
                f"code {srvproc.returncode}"
            )
        print(f"- sending a SIGTERM signal to the {proto}server process")
        try:
            srvproc.terminate()
        except OSError as err:
            sys.exit(
                f"Could not send a SIGTERM signal to the {proto}server process at "
                f"{srvproc.pid}: {err}"
            )
        print(f"- waiting for the {proto}server process to exit")
        res = srvproc.wait()
        if res != 0:
            sys.exit(f"The {proto}server process exited with a non-zero code {res}")

    def wait_and_connect(srvproc: subprocess.Popen[str]) -> None:
        """Wait for the "ready" message from tcpserver."""
        assert srvproc.stderr is not None, repr(srvproc)
        print(f"- awaiting the first 'status' line from {proto}server")
        line = srvproc.stderr.readline()
        if "server: status: 0/" not in line:
            sys.exit(
                f"Unexpected first line from {proto}server: expected 'status: 0/N', got {line!r}"
            )
        with_child(
            runner,
            "test_server_client_spew",
            [str(clipath), "-R", "--"] + addr + ["sh", "-c", "set -e; exec <&6; exec cat"],
            lambda cliproc: run(srvproc, cliproc),
        )

    proto = runner.proto
    print(f"\n=== Testing {proto}server against {proto}client")

    srvpath = runner.cfg.bindir / f"{proto}server"
    print(f"- will spawn {proto}server at {srvpath} in a while")
    clipath = runner.cfg.bindir / f"{proto}client"
    print(f"- will spawn {proto}client at {clipath} in a while")
    message = MSG_RESP_HELLO + MSG_RESP_BYE
    with_child(
        runner,
        "test_server_client_spew",
        [str(srvpath), "-vR", "--"] + addr + ["printf", "--", message.replace("\n", "\\n")],
        wait_and_connect,
        stderr=subprocess.PIPE,
    )

    print(f"- {srvpath} seems fine")


def run_test(runner: Runner) -> None:
    """Run a couple of UCSPI tests."""
    addr = runner.find_listening_address()

    test_local_cat(runner, addr)
    test_local_client_spew(runner, addr)

    test_server_local(runner, addr)
    test_server_client_spew(runner, addr)

    print(f"\n=== The tests for {runner.cfg.proto} passed")


def add_handler(proto: str, runner: type[Runner]) -> None:
    """Add a UCSPI protocol test runner."""
    current = HANDLERS.get(proto)
    if current is None:
        HANDLERS[proto] = runner
    elif current != runner:
        raise RunnerError(
            f"Handler mismatch for the {proto!r} protocol: had {current!r}, now {runner!r}"
        )


def run_test_handler(cfg: Config) -> None:
    """Parse command-line arguments, run the tests."""
    hprot = HANDLERS.get(cfg.proto)
    if hprot is None:
        sys.exit(f"Don't know how to test the {cfg.proto!r} UCSPI protocol")

    run_test(hprot(cfg, cfg.proto))
