"""Does listen.udp.workers help or hurt, on this platform, with this tracker?

Starts a tracker once per worker count and drives it with tests/udp_load.c,
which speaks the real protocol, so what comes out is the tracker's own
throughput rather than an echo server's.

    cc -O2 -pthread -o /tmp/udp_load tests/udp_load.c
    make opentracker
    python3 tests/udp_scaling.py

Worker count 0 means no dedicated workers: udp is served from the main event
loop, which is what a default build does.

LOAD_THREADS and WINDOW decide how much load is offered. Both matter. Too
little and the tracker idles while the figure measures round trips; too much
and the kernel drops datagrams, every lost reply costs the generator a
timeout, and the figure collapses. The generator prints how many were lost:
tune until that is near zero and the throughput has stopped rising.
"""
import os
import signal
import socket
import subprocess
import sys
import tempfile
import time

TRACKER = os.environ.get("OPENTRACKER", "./opentracker")
LOADER = os.environ.get("UDP_LOAD", "/tmp/udp_load")
WORKER_COUNTS = [int(n) for n in os.environ.get("WORKERS", "0 1 2 4 8").split()]
LOAD_THREADS = int(os.environ.get("LOAD_THREADS", "4"))
SECONDS = float(os.environ.get("SECONDS", "5"))
TORRENTS = int(os.environ.get("TORRENTS", "10000"))
WINDOW = int(os.environ.get("WINDOW", "32"))
REPEATS = int(os.environ.get("REPEATS", "3"))
PORT_BASE = int(os.environ.get("PORT_BASE", "19960"))


def start(port, workers):
    directory = tempfile.mkdtemp(prefix="ot-udpscale-")
    config = "listen.tcp 127.0.0.1:%d\nlisten.udp 127.0.0.1:%d\n" % (port, port)
    if workers:
        config += "listen.udp.workers %d\n" % workers
    with open(os.path.join(directory, "conf"), "w") as handle:
        handle.write(config)
    log = open(os.path.join(directory, "log"), "w+b")
    tracker = subprocess.Popen([TRACKER, "-f", os.path.join(directory, "conf")], stdout=log, stderr=log)
    for _ in range(300):
        if tracker.poll() is not None:
            log.seek(0)
            raise SystemExit("tracker exited at once on port %d with %d workers:\n%s"
                             % (port, workers, log.read().decode("utf-8", "replace")))
        try:
            socket.create_connection(("127.0.0.1", port), timeout=0.2).close()
            return tracker, log
        except OSError:
            time.sleep(0.02)
    tracker.kill()
    raise SystemExit("tracker never listened on %d" % port)


def cpu_seconds(pid):
    """Cumulative cpu time of a process, the one spelling ps has everywhere."""
    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, parts = 0.0, out.replace("-", ":").split(":")
    for part in parts:
        total = total * 60 + float(part)
    return total


def stop(tracker):
    tracker.send_signal(signal.SIGINT)
    try:
        tracker.wait(timeout=10)
    except subprocess.TimeoutExpired:
        tracker.kill()
        tracker.wait()


def main():
    if not os.path.exists(LOADER):
        raise SystemExit("build the generator first: cc -O2 -pthread -o %s tests/udp_load.c" % LOADER)

    print("  tracker %s, %d load threads x %d in flight, %d torrents, %.0f s x %d runs"
          % (TRACKER, LOAD_THREADS, WINDOW, TORRENTS, SECONDS, REPEATS))
    print("  %-8s %12s %12s %12s %7s   %s"
          % ("workers", "median/s", "lowest", "highest", "cores", "vs 0 workers"))
    baseline = None
    for index, workers in enumerate(WORKER_COUNTS):
        rates, cores = [], []
        for repeat in range(REPEATS):
            port = PORT_BASE + index * 16 + repeat
            tracker, log = start(port, workers)
            try:
                # Warm up until the torrent set exists and the peers are in
                # place, so the measured run is renewals rather than insertions
                subprocess.run([LOADER, "127.0.0.1", str(port), str(LOAD_THREADS), "2", str(TORRENTS), str(WINDOW)],
                               stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
                spent, wall = cpu_seconds(tracker.pid), time.time()
                result = subprocess.run(
                    [LOADER, "127.0.0.1", str(port), str(LOAD_THREADS), str(SECONDS), str(TORRENTS), str(WINDOW)],
                    capture_output=True, text=True,
                )
                used = cpu_seconds(tracker.pid)
            finally:
                stop(tracker)
                log.close()
            head = result.stdout.strip().split()
            if head and head[0].replace(".", "").isdigit():
                rates.append(float(head[0]))
                if used is not None and spent is not None:
                    cores.append((used - spent) / max(time.time() - wall, 1e-9))

        if not rates:
            print("  %-8s no result" % workers)
            continue
        rates.sort()
        median = rates[len(rates) // 2]
        if baseline is None:
            baseline = median
        # A spread wider than the differences between rows means the rows are
        # not telling you anything
        busy = ("%7.2f" % (sorted(cores)[len(cores) // 2])) if cores else "      ?"
        print("  %-8s %12.0f %12.0f %12.0f %s   %.2fx  (spread %.0f%%)"
              % (workers, median, rates[0], rates[-1], busy, median / baseline,
                 100.0 * (rates[-1] - rates[0]) / median))


if __name__ == "__main__":
    main()