#!/usr/bin/env python3
"""Field test without any sensor hardware.

  udpprobe.py listen PORT [PORT ...]        print every datagram arriving on these ports
  udpprobe.py send PORT [--to 192.168.1.255] [--rate 10] [--count 0]

  Round trip on ONE clock (true latency, immune to clock differences):
  udpprobe.py echo IN_PORT OUT_PORT [--to 192.168.1.255]      on the far site: bounce probes back
  udpprobe.py ping OUT_PORT IN_PORT [--to 192.168.71.255] [--count 50] [--rate 10]

  Example (Essen pings Berlin; 4999 is a Folkwang port, 8999 a UdK port):
    Berlin:  python3 udpprobe.py echo 4999 8999 --to 192.168.1.255
    Essen:   python3 udpprobe.py ping 4999 8999 --to 192.168.71.255

`send` puts the wall-clock time in each packet; `listen` on the other site prints
"one-way" = receiver clock - sender clock, which includes any clock offset between
the two machines. Use echo/ping for real numbers: one-way ~= rtt / 2.
"""
import socket, sys, time, selectors, struct

def sock(bind_port=None):
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    if hasattr(socket, "SO_REUSEPORT"):
        try: s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
        except OSError: pass
    if bind_port is not None:
        s.bind(("0.0.0.0", bind_port))
    return s

def listen(ports):
    sel = selectors.DefaultSelector()
    for p in ports:
        sel.register(sock(p), selectors.EVENT_READ, p)
    print("listening on", ports)
    while True:
        for k, _ in sel.select():
            d, src = k.fileobj.recvfrom(65535)
            extra = ""
            if d.startswith(b"PROBE") and len(d) >= 17:
                seq, ts = struct.unpack("!Id", d[5:17])
                extra = f" seq={seq} one-way={(time.time()-ts)*1000:.1f}ms (incl. clock offset)"
            print(f"{time.strftime('%H:%M:%S')} port {k.data} from {src[0]}:{src[1]} len {len(d)}{extra}", flush=True)

def send(port, to, rate, count):
    s = sock()
    i = 0
    while count == 0 or i < count:
        s.sendto(b"PROBE" + struct.pack("!Id", i, time.time()) + b"....", (to, port))
        i += 1; time.sleep(1 / rate)

def echo(in_port, out_port, to):
    r, w = sock(in_port), sock()
    print(f"echoing probes from port {in_port} back to {to}:{out_port}")
    n = 0
    while True:
        d, src = r.recvfrom(65535)
        if d.startswith(b"PING"):
            w.sendto(b"PONG" + d[4:], (to, out_port))
            n += 1
            if n % 10 == 1: print(f"{time.strftime('%H:%M:%S')} echoed {n} probes", flush=True)

def ping(out_port, in_port, to, count, rate):
    r, w = sock(in_port), sock()
    r.settimeout(0.001)
    sent, rtts = {}, []
    t_end = None; i = 0; nxt = time.perf_counter()
    while True:
        now = time.perf_counter()
        if i < count and now >= nxt:
            w.sendto(b"PING" + struct.pack("!Id", i, now), (to, out_port))
            sent[i] = now; i += 1; nxt += 1 / rate
            if i == count: t_end = now + 2.0
        try:
            d, _ = r.recvfrom(65535)
            if d.startswith(b"PONG") and len(d) >= 16:
                seq, ts = struct.unpack("!Id", d[4:16])
                if seq in sent:
                    rtt = (time.perf_counter() - sent.pop(seq)) * 1000
                    rtts.append(rtt)
                    print(f"seq={seq} rtt={rtt:.1f} ms  (one-way ~{rtt/2:.1f} ms)", flush=True)
        except socket.timeout:
            pass
        if t_end and time.perf_counter() > t_end: break
    if rtts:
        rtts.sort()
        q = lambda f: rtts[min(len(rtts) - 1, int(f * len(rtts)))]
        print(f"\n{len(rtts)}/{count} answered  rtt min/median/p95/max = "
              f"{rtts[0]:.1f}/{q(.5):.1f}/{q(.95):.1f}/{rtts[-1]:.1f} ms  -> one-way ~{q(.5)/2:.1f} ms")
    else:
        print(f"\n0/{count} answered - is `echo` running on the other site?")

if __name__ == "__main__":
    a = sys.argv[1:]
    if not a: sys.exit(__doc__)
    opt = lambda k, d: a[a.index(k) + 1] if k in a else d
    if a[0] == "listen": listen([int(x) for x in a[1:] if x.isdigit()])
    elif a[0] == "send":
        send(int(a[1]), opt("--to", "192.168.1.255"), float(opt("--rate", "10")), int(opt("--count", "0")))
    elif a[0] == "echo":
        echo(int(a[1]), int(a[2]), opt("--to", "192.168.1.255"))
    elif a[0] == "ping":
        ping(int(a[1]), int(a[2]), opt("--to", "192.168.71.255"), int(opt("--count", "50")), float(opt("--rate", "10")))
    else: sys.exit(__doc__)