#!/usr/bin/env python3
"""Test a paced stick-table handover between old and new local workers.

The test uses one HAProxy master, loopback sockets, and documentation-only IPv6
keys. It rewrites the configuration before a soft reload so the old worker's
connection to the new worker passes through a deliberately slow TCP relay.
This makes transfers exceed HAProxy's historical five-second peer stream
timeout using only loopback sockets.
"""

import argparse
import os
from pathlib import Path
import re
import signal
import socket
import subprocess
import tempfile
import threading
import time


def command(path: Path, request: str) -> str:
    """Send a runtime command and return its response.

    Args:
        path: Unix socket accepting the command.
        request: HAProxy Runtime API command without its trailing newline.

    Returns:
        The complete decoded response.

    Raises:
        OSError: The socket cannot be reached or read.
    """
    with socket.socket(socket.AF_UNIX) as client:
        client.settimeout(10)
        client.connect(str(path))
        client.sendall((request + "\n").encode())
        chunks: list[bytes] = []
        while chunk := client.recv(65536):
            chunks.append(chunk)
        return b"".join(chunks).decode()


def listener() -> socket.socket:
    """Return a listening socket bound to an available loopback TCP port."""
    server = socket.socket()
    server.bind(("127.0.0.1", 0))
    server.listen()
    return server


def key(index: int) -> str:
    """Return a unique documentation-only IPv6 address for index."""
    return f"2001:db8::{index >> 16:x}:{index & 65535:x}"


class Relay:
    """Forward TCP connections while pacing each direction independently."""

    def __init__(self, server: socket.socket, destination: int, delay: float):
        """Start a relay from server to destination.

        Args:
            server: Bound listener owned by the relay.
            destination: Loopback TCP port receiving forwarded connections.
            delay: Seconds to wait after forwarding each 1 KiB chunk.
        """
        self.server = server
        self.destination = destination
        self.delay = delay
        self.stopped = threading.Event()
        self.lock = threading.Lock()
        self.sockets: set[socket.socket] = set()
        self.workers: list[threading.Thread] = []
        self.thread = threading.Thread(target=self._accept, daemon=True)
        self.thread.start()

    def _accept(self) -> None:
        """Accept connections and start workers for both directions."""
        self.server.settimeout(0.2)
        while not self.stopped.is_set():
            try:
                client, _ = self.server.accept()
            except TimeoutError:
                continue
            except OSError:
                break
            try:
                remote = socket.create_connection(
                    ("127.0.0.1", self.destination), timeout=1,
                )
                remote.settimeout(None)
            except OSError:
                client.close()
                continue
            with self.lock:
                self.sockets.update((client, remote))
            for source, target in ((client, remote), (remote, client)):
                worker = threading.Thread(
                    target=self._forward,
                    args=(source, target),
                    daemon=True,
                )
                self.workers.append(worker)
                worker.start()

    def _forward(self, source: socket.socket, target: socket.socket) -> None:
        """Copy paced chunks until either endpoint closes or the relay stops."""
        try:
            while not self.stopped.is_set():
                chunk = source.recv(1024)
                if not chunk:
                    break
                target.sendall(chunk)
                self.stopped.wait(self.delay)
        except OSError:
            pass
        finally:
            for connection in (source, target):
                try:
                    connection.shutdown(socket.SHUT_RDWR)
                except OSError:
                    pass
                connection.close()
                with self.lock:
                    self.sockets.discard(connection)

    def close(self) -> None:
        """Stop forwarding and close all relay-owned sockets."""
        self.stopped.set()
        self.server.close()
        self.thread.join(timeout=2)
        with self.lock:
            connections = list(self.sockets)
        for connection in connections:
            try:
                connection.shutdown(socket.SHUT_RDWR)
            except OSError:
                pass
            connection.close()
        for worker in self.workers:
            worker.join(timeout=2)


def write_config(
    path: Path,
    admin: Path,
    peer_port: int,
    entries: int,
    threads: int,
    incoming_tls: str = "",
    outgoing_tls: str = "",
    transport: tuple[int, int] | None = None,
) -> None:
    """Write the initial or reload configuration.

    Args:
        path: Destination configuration path.
        admin: Runtime API socket path.
        peer_port: Local peer listener address.
        entries: Stick-table entries required by the test.
        threads: Worker thread count.
        incoming_tls: Optional TLS arguments for the peer listener.
        outgoing_tls: Optional TLS arguments for the local peer server.
        transport: Optional old-listener and relay-destination port pair.
    """
    config = f"""global
  localpeer lab
  nbthread {threads}
  stats socket {admin} level admin
peers synced
  bind 127.0.0.1:{peer_port}{incoming_tls}
  server lab{outgoing_tls}
  table marker type ipv6 size {entries + 100} expire 1h store gpt(1)
"""
    if transport:
        old_port, relay_port = transport
        config += f"""listen handover-transport
  mode tcp
  bind 127.0.0.1:{old_port}
  timeout connect 1s
  timeout client 1m
  timeout server 1m
  server paced 127.0.0.1:{relay_port}
"""
    path.write_text(config)


def wait_for_ready(admin: Path, process: subprocess.Popen[str]) -> int:
    """Wait for initial peer synchronization and return the worker PID."""
    deadline = time.monotonic() + 15
    while time.monotonic() < deadline:
        if process.poll() is not None:
            raise RuntimeError("HAProxy exited before becoming ready")
        try:
            peers = command(admin, "show peers")
            info = command(admin, "show info")
        except OSError:
            time.sleep(0.05)
            continue
        flags = re.search(r"disabled=0 flags=(0x[0-9a-f]+)", peers)
        pid = re.search(r"^Pid: (\d+)", info, re.MULTILINE)
        if flags and int(flags[1], 16) & 3 == 3 and pid:
            return int(pid[1])
        time.sleep(0.05)
    raise RuntimeError("HAProxy did not finish initial peer synchronization")


def seed(admin: Path, entries: int) -> None:
    """Insert entries documentation-only markers into the local table."""
    for offset in range(0, entries, 100):
        request = ";".join(
            f"set table synced/marker key {key(index)} data.gpt[0] 93"
            for index in range(offset, min(offset + 100, entries))
        )
        reply = command(admin, request)
        if reply.strip():
            raise RuntimeError(f"Cannot seed stick table: {reply[:200]}")


def table_count(admin: Path) -> int:
    """Return the number of entries in the synthetic table."""
    output = command(admin, "show table")
    match = re.search(r"# table: synced/marker,.*used:(\d+)", output)
    if not match:
        raise RuntimeError(f"Synthetic table is missing: {output}")
    return int(match[1])


def check(binary: Path, entries: int, threads: int, delay: float,
          timeout: float, tls: bool) -> bool:
    """Run one paced local handover and return whether it completes.

    Args:
        binary: HAProxy executable under test.
        entries: Number of synthetic entries to transfer.
        threads: Worker thread count.
        delay: Relay delay after each 1 KiB chunk.
        timeout: Maximum seconds allowed for the handover.
        tls: Whether to encrypt and mutually authenticate the peer stream.

    Returns:
        True when all entries arrive and synchronization finishes.
    """
    with tempfile.TemporaryDirectory(prefix="ha-local-handover-") as directory:
        root = Path(directory)
        admin = root / "admin.sock"
        config = root / "haproxy.cfg"
        old_listener = listener()
        old_port = int(old_listener.getsockname()[1])
        old_listener.close()
        incoming_tls = outgoing_tls = ""
        if tls:
            key_path = root / "key.pem"
            certificate_path = root / "certificate.pem"
            bundle_path = root / "peer.pem"
            subprocess.run(
                [
                    "openssl", "req", "-x509", "-newkey", "rsa:2048", "-nodes",
                    "-keyout", str(key_path), "-out", str(certificate_path),
                    "-days", "1", "-subj", "/CN=lab",
                ],
                check=True,
                capture_output=True,
            )
            bundle_path.write_bytes(
                certificate_path.read_bytes() + key_path.read_bytes(),
            )
            incoming_tls = (
                f" ssl crt {bundle_path} verify required "
                f"ca-file {certificate_path}"
            )
            outgoing_tls = (
                f" ssl verify required ca-file {certificate_path} "
                f"crt {bundle_path}"
            )
        write_config(
            config,
            admin,
            old_port,
            entries,
            threads,
            incoming_tls,
            outgoing_tls,
        )
        log_path = root / "haproxy.log"
        with log_path.open("w") as log:
            process = subprocess.Popen(
                [str(binary), "-W", "-db", "-f", str(config)],
                stdout=log,
                stderr=log,
                text=True,
                start_new_session=True,
            )
        relay: Relay | None = None
        try:
            old_pid = wait_for_ready(admin, process)
            seed(admin, entries)
            if table_count(admin) != entries:
                raise RuntimeError("Source table was not fully seeded")

            new_listener = listener()
            new_port = int(new_listener.getsockname()[1])
            new_listener.close()
            relay_listener = listener()
            relay_port = int(relay_listener.getsockname()[1])
            relay = Relay(relay_listener, new_port, delay)
            write_config(
                config,
                admin,
                new_port,
                entries,
                threads,
                incoming_tls,
                outgoing_tls,
                (old_port, relay_port),
            )

            os.kill(process.pid, signal.SIGUSR2)
            started = time.monotonic()
            while time.monotonic() - started < 15:
                try:
                    info = command(admin, "show info")
                    pid = re.search(r"^Pid: (\d+)", info, re.MULTILINE)
                    if pid and int(pid[1]) != old_pid:
                        break
                except OSError:
                    pass
                time.sleep(0.05)
            else:
                raise RuntimeError("Reloaded worker did not become ready")

            last_count = 0
            last_flags = "missing"
            while time.monotonic() - started < timeout:
                last_count = table_count(admin)
                peers = command(admin, "show peers")
                flags = re.search(r"disabled=0 flags=(0x[0-9a-f]+)", peers)
                last_flags = flags[1] if flags else "missing"
                if (last_count == entries and flags and
                        int(last_flags, 16) & 3 == 3):
                    elapsed = time.monotonic() - started
                    print(
                        f"complete=true entries={last_count} "
                        f"seconds={elapsed:.2f} flags={last_flags}",
                    )
                    return True
                time.sleep(0.25)
            print(
                f"complete=false entries={last_count} "
                f"seconds={time.monotonic() - started:.2f} "
                f"flags={last_flags}",
            )
            return False
        finally:
            if relay:
                relay.close()
            if process.poll() is None:
                os.killpg(process.pid, signal.SIGTERM)
                process.wait(timeout=10)


def main() -> None:
    """Parse arguments and exit nonzero when local handover fails."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("binary", type=Path)
    parser.add_argument("--entries", type=int, default=100_000)
    parser.add_argument("--threads", type=int, default=8)
    parser.add_argument("--delay", type=float, default=0.005)
    parser.add_argument("--timeout", type=float, default=30)
    parser.add_argument("--tls", action="store_true")
    args = parser.parse_args()
    if args.entries < 1 or args.threads < 1:
        parser.error("entries and threads must be positive")
    if args.delay < 0 or args.timeout <= 0:
        parser.error("delay must be nonnegative and timeout must be positive")
    passed = check(
        args.binary.resolve(),
        args.entries,
        args.threads,
        args.delay,
        args.timeout,
        args.tls,
    )
    raise SystemExit(0 if passed else 1)


if __name__ == "__main__":
    main()