"""The four udp arrangements, side by side.

    cc -O2 -pthread -o /tmp/udp_modes tests/udp_modes.c
    cc -O2 -pthread -o /tmp/udp_load  tests/udp_load.c
    python3 tests/udp_modes.py

Two questions at once. Down the table: does a shared socket stop scaling as
threads are added? Across it: does taking several datagrams per syscall make
that better, or only cheaper?

The receiver does no tracker work by default, so the figure is the kernel's
receive path and nothing else. WORK_NS=1500 dilutes it back towards what the
tracker actually costs per announce.

The number reported is cpu microseconds the receiver spent per datagram
answered, not throughput. Throughput here says more about the load generator
than the receiver, because both run on the same machine and an echo costs so
little that the generator runs out of cores first. Cpu per datagram does not
care: it is what the receive path costs, whoever is the bottleneck.
"""
import os
import signal
import socket
import subprocess
import sys
import time

MODES = os.environ.get("MODES", "shared per-thread").split()
BATCHES = [int(n) for n in os.environ.get("BATCHES", "1 8 32").split()]
THREAD_COUNTS = [int(n) for n in os.environ.get("THREADS", "1 2 4 8").split()]
LOAD_THREADS = int(os.environ.get("LOAD_THREADS", "4"))
WINDOW = int(os.environ.get("WINDOW", "32"))
SECONDS = float(os.environ.get("SECONDS", "3"))
REPEATS = int(os.environ.get("REPEATS", "2"))
WORK_NS = int(os.environ.get("WORK_NS", "0"))
RECEIVER = os.environ.get("UDP_MODES", "/tmp/udp_modes")
LOADER = os.environ.get("UDP_LOAD", "/tmp/udp_load")
PORT_BASE = int(os.environ.get("PORT_BASE", "20100"))


def cpu_seconds(pid):
    try:
        out = subprocess.run(["ps", "-o", "time=", "-p", str(pid)], capture_output=True, text=True).stdout.strip()
    except OSError:
        return None
    if not out:
        return None
    total = 0.0
    for part in out.replace("-", ":").split(":"):
        total = total * 60 + float(part)
    return total


def one_run(port, threads, mode, batch):
    receiver = subprocess.Popen(
        [RECEIVER, str(port), str(threads), mode, str(batch), str(WORK_NS)],
        stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True,
    )
    for _ in range(200):
        if receiver.poll() is not None:
            return None, None
        probe = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        probe.settimeout(0.1)
        try:
            probe.sendto(b"\x00" * 16, ("127.0.0.1", port))
            probe.recv(64)
            probe.close()
            break
        except OSError:
            probe.close()
            time.sleep(0.02)
    else:
        receiver.send_signal(signal.SIGINT)
        receiver.wait()
        return None, None

    try:
        subprocess.run([LOADER, "127.0.0.1", str(port), str(LOAD_THREADS), "1", "10000", str(WINDOW)],
                       stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        spent, wall = cpu_seconds(receiver.pid), time.time()
        out = subprocess.run(
            [LOADER, "127.0.0.1", str(port), str(LOAD_THREADS), str(SECONDS), "10000", str(WINDOW)],
            capture_output=True, text=True,
        ).stdout.strip()
        used = cpu_seconds(receiver.pid)
    finally:
        receiver.send_signal(signal.SIGINT)
        try:
            receiver.wait(timeout=5)
        except subprocess.TimeoutExpired:
            receiver.kill()
            receiver.wait()

    head = out.split()
    if not head or not head[0].replace(".", "").isdigit():
        return None, None
    cores = (used - spent) / max(time.time() - wall, 1e-9) if (used is not None and spent is not None) else None
    return float(head[0]), cores


def main():
    for binary in (RECEIVER, LOADER):
        if not os.path.exists(binary):
            raise SystemExit("missing %s; see the comment at the top of tests/udp_modes.c" % binary)

    print("  %d load threads x %d in flight, %.0f s x %d runs, %d ns work per datagram"
          % (LOAD_THREADS, WINDOW, SECONDS, REPEATS, WORK_NS))
    print("  cpu microseconds per datagram, and the throughput it was measured at")
    header = "  %-14s %-8s" % ("sockets", "threads")
    for batch in BATCHES:
        header += "%20s" % ("%d per syscall" % batch)
    print(header)

    port = PORT_BASE
    for mode in MODES:
        for threads in THREAD_COUNTS:
            row = "  %-14s %-8d" % (mode, threads)
            for batch in BATCHES:
                results = []
                for _ in range(REPEATS):
                    port += 1
                    rate, cores = one_run(port, threads, mode, batch)
                    if rate:
                        results.append((rate, cores))
                if not results:
                    row += "%18s" % "-"
                    continue
                results.sort()
                rate, cores = results[len(results) // 2]
                if cores:
                    row += "%9.2f us %6.0fk/s" % (cores * 1e6 / rate, rate / 1000.0)
                else:
                    row += "%20s" % ("%.0fk/s" % (rate / 1000.0))
            print(row)
            sys.stdout.flush()


if __name__ == "__main__":
    main()