#!/usr/bin/env python3
"""Test peer heartbeats during slow full resynchronization with local writes.

Uses two isolated HAProxy processes and paced loopback relays. No reload is
involved. A successful run must learn every seeded entry,
retain the receiving node's marker, and replicate its deferred write back.
Requires Python 3.13, loopback socket access, and openssl for optional mTLS.
"""

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


MARKER = "2001:db8:ffff::1"


def command(path: Path, request: str) -> str:
    """Send request to a runtime socket; return its reply or raise on I/O errors."""
    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:
    """Reserve and return an IPv4 loopback listening socket on a free port."""
    server = socket.socket()
    server.bind(("127.0.0.1", 0))
    server.listen()
    return server


class Relay:
    """Pace both directions of loopback TCP sessions without inspecting payloads.

    The caller owns the listener. close() stops accepting, closes active sockets,
    and joins worker threads. Failed connections are closed, not retried.
    """

    def __init__(self, server: socket.socket, destination: int, delay: float):
        """Start forwarding server to destination, delaying each 1 KiB by seconds."""
        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 loopback connections and start one worker for each direction."""
        self.server.settimeout(0.2)
        while not self.stopped.is_set():
            try:
                client, _ = self.server.accept()
            except socket.timeout:
                continue
            except OSError:
                break
            try:
                remote = socket.create_connection(("127.0.0.1", self.destination), 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, closing both directions on EOF or a socket error."""
        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 the relay and release its listener, sockets, and worker threads."""
        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 key(index: int) -> str:
    """Return a unique documentation-only IPv6 key for a nonnegative index."""
    return f"2001:db8::{index >> 16:x}:{index & 65535:x}"


def run(binary: Path, entries: int, threads: int, delay: float,
        timeout: int, tls: bool, idle: bool) -> bool:
    """Exercise binary with paced peer traffic; return whether resync succeeds.

    entries controls the initial table size; threads sets each node's workers;
    delay is seconds per relay chunk; timeout bounds learning after node B
    starts. tls enables mutual TLS; idle disables writes during learning for
    a control experiment. Temporary files and owned processes are always cleaned.
    Raises on startup, configuration, runtime socket, or subprocess errors.
    """
    with tempfile.TemporaryDirectory(prefix="ha-resync-") as directory:
        root = Path(directory)
        reserved = [listener() for _ in range(4)]
        ports = [server.getsockname()[1] for server in reserved]
        processes: list[subprocess.Popen[str]] = []
        relays: list[Relay] = []
        sockets = [root / "a.sock", root / "b.sock"]
        try:
            incoming = outgoing = ""
            if tls:
                subprocess.run([
                    "openssl", "req", "-x509", "-newkey", "rsa:2048", "-nodes",
                    "-keyout", str(root / "key.pem"), "-out", str(root / "cert.pem"),
                    "-days", "1", "-subj", "/CN=lab",
                ], check=True, capture_output=True)
                (root / "node.pem").write_bytes(
                    (root / "cert.pem").read_bytes() + (root / "key.pem").read_bytes(),
                )
                incoming = f" ssl crt {root}/node.pem verify required ca-file {root}/cert.pem"
                outgoing = f"  default-server ssl verify required ca-file {root}/cert.pem crt {root}/node.pem\n"
            for index in range(2):
                relays.append(Relay(reserved[index + 2], ports[index], delay))
            for index, node in enumerate(("a", "b")):
                config = root / f"{node}.cfg"
                config.write_text(f"""global
  localpeer {node}
  nbthread {threads}
  stats socket {sockets[index]} level admin
peers synced
  bind 127.0.0.1:{ports[index]}{incoming}
{outgoing}  server {node}
  server {'b' if node == 'a' else 'a'} 127.0.0.1:{ports[3 - index]}
  table guard type ipv6 size {entries + 100} expire 1h store gpt(1),http_req_rate(1m),http_err_rate(1m)
""")
                reserved[index].close()
                with (root / f"{node}.log").open("w") as log:
                    process = subprocess.Popen(
                        [str(binary), "-db", "-f", str(config)],
                        stdout=log, stderr=log, text=True, start_new_session=True,
                    )
                processes.append(process)
                deadline = time.monotonic() + 10
                while not sockets[index].exists():
                    if process.poll() is not None or time.monotonic() > deadline:
                        raise RuntimeError((root / f"{node}.log").read_text())
                    time.sleep(0.05)
                command(sockets[index], "trace peers event none event sess_end event sess_shut event sess_err event proto_error event sess_resync level state sink buf0 start now")
                if index == 0:
                    for offset in range(0, entries, 100):
                        request = ";".join(
                            f"set table synced/guard key {key(i)} data.gpt[0] 93"
                            for i in range(offset, min(offset + 100, entries))
                        )
                        reply = command(sockets[0], request)
                        if reply.strip():
                            raise RuntimeError(f"Seeding failed: {reply[:200]}")
                    # Wait for the source to finish its initial sync timeout.
                    deadline = time.monotonic() + 15
                    while True:
                        flags = re.search(r"disabled=0 flags=(0x[0-9a-f]+)", command(sockets[0], "show peers"))
                        if flags and int(flags[1], 16) & 3 == 3:
                            break
                        if time.monotonic() > deadline:
                            raise RuntimeError("Source did not finish initial sync")
                        time.sleep(0.1)
            started = time.monotonic()
            last_write = 0
            observed_learning = False
            while time.monotonic() - started < timeout:
                time.sleep(1)
                if not idle:
                    last_write += 1
                    command(sockets[1], f"set table synced/guard key {MARKER} data.gpt[0] {last_write}")
                peers = command(sockets[1], "show peers")
                observed_learning |= "learn_state=PROCESSING" in peers
                flags = re.search(r"disabled=0 flags=(0x[0-9a-f]+)", peers)
                used = re.search(r"used:(\d+)", command(sockets[1], "show table"))
                summary = [line.strip() for line in peers.splitlines() if "(remote," in line or "tx_hbt=" in line]
                elapsed = time.monotonic() - started
                print(f"elapsed={elapsed:.1f}s used={used[1] if used else '?'} " + " ".join(summary), flush=True)
                if flags and int(flags[1], 16) & 3 == 3 and used and int(used[1]) == entries + (not idle):
                    if not observed_learning or elapsed < 8:
                        raise RuntimeError("Resync was too fast; increase entries or relay delay")
                    for sample in (0, entries // 2, entries - 1):
                        if "gpt0=93" not in command(sockets[1], f"show table synced/guard key {key(sample)}"):
                            raise RuntimeError("Learned sample value is incorrect")
                    deadline = time.monotonic() + 10
                    while not idle:
                        response = command(sockets[0], f"show table synced/guard key {MARKER}")
                        value = re.search(r"gpt0=(\d+)", response)
                        if value and int(value[1]) == last_write:
                            break
                        if time.monotonic() > deadline:
                            raise RuntimeError("Deferred local write did not reach source")
                        time.sleep(0.1)
                    print("PASS: full resync completed; sampled values verified"
                          + ("" if idle else "; deferred write verified"), flush=True)
                    return True
            print("FAIL: full resync did not finish", flush=True)
            for index, node in enumerate(("a", "b")):
                print(f"--- {node} ---\n" + command(sockets[index], "show peers"))
                print(command(sockets[index], "show events buf0"))
            return False
        finally:
            for process in processes:
                if process.poll() is None:
                    os.killpg(process.pid, signal.SIGTERM)
                    process.wait(timeout=5)
            for relay in relays:
                relay.close()
            for server in reserved:
                server.close()


def main() -> None:
    """Parse test options and exit nonzero on failed synchronization or errors."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("binary", type=Path)
    parser.add_argument("--entries", type=int, default=100000)
    parser.add_argument("--threads", type=int, default=1)
    parser.add_argument("--delay", type=float, default=0.005)
    parser.add_argument("--timeout", type=int, default=45)
    parser.add_argument("--tls", action="store_true")
    parser.add_argument("--idle", action="store_true")
    args = parser.parse_args()
    if args.entries < 1000 or args.threads < 1 or args.delay <= 0 or args.timeout < 10:
        parser.error("entries >= 1000, threads >= 1, delay > 0, timeout >= 10 required")
    raise SystemExit(0 if run(args.binary.resolve(), args.entries, args.threads,
                            args.delay, args.timeout, args.tls, args.idle) else 1)


if __name__ == "__main__":
    main()