#!/usr/bin/env python3
"""Orin side of the app's link test (lib/linktest.dart). Development tool, not a service.

The tablet's own sockets are the only ones bound to the USB link (`adb shell`
traffic goes over WiFi), so the app drives the test and this end only serves:

  7791/tcp  down: client sends one line, then reads until the server closes:
              "<id> <secs>"               zeros for <secs>
              "<id> burst"                then per line "<bytes>": exactly <bytes> zeros,
                                          on the same connection until the client closes
              "<id> paced <B/s> <secs>"   30 frames/s, each [u32 length][f64 send time][zeros],
                                          length counting the whole frame
  7792/tcp  up:   client sends "<id>\\n" then data until EOF; gets "<bytes>\\n" back
  7793/udp  echo
  7794/tcp  results: one JSON line per connection, appended to the results file

Every TCP test also logs this end's TCP_INFO (retransmits, RTT), which the app
cannot read. Run: python3 linktest_server.py [results.jsonl]

LINKTEST_PACE=<bytes/s> caps every download socket with SO_MAX_PACING_RATE, TCP's
own pacing (no qdisc needed; the Tegra kernel ships without fq/tbf/htb).
"""

import json
import os
import socket
import struct
import sys
import threading
import time

PACE = int(os.environ.get("LINKTEST_PACE", "0"))
SO_MAX_PACING_RATE = 47
RESULTS = os.path.expanduser(sys.argv[1] if len(sys.argv) > 1 else "~/linktest/results.jsonl")
CHUNK = b"\0" * 65536
lock = threading.Lock()


def log(rec):
    rec = {"orin_ts": round(time.time(), 3), **rec}
    with lock, open(RESULTS, "a") as f:
        f.write(json.dumps(rec) + "\n")
    print(json.dumps(rec), flush=True)


def tcp_info(sock):
    # struct tcp_info: 8 bytes of u8 fields, then u32s; rtt is the 16th, total_retrans the 24th.
    raw = sock.getsockopt(socket.IPPROTO_TCP, socket.TCP_INFO, 104)
    u = struct.unpack_from("8B24I", raw)
    return {"retrans": u[8 + 23], "rtt_ms": u[8 + 15] / 1000, "rto_ms": u[8] / 1000}


def readline(conn):
    buf = b""
    while not buf.endswith(b"\n") and len(buf) < 256:
        b = conn.recv(1)
        if not b:
            break
        buf += b
    return buf.decode().strip()


def down(conn, peer):
    test_id, *args = readline(conn).split()
    mode = args[0] if args[0] in ("burst", "paced") else "bulk"
    if PACE:
        conn.setsockopt(socket.SOL_SOCKET, SO_MAX_PACING_RATE, PACE)
    sent = 0
    try:
        if mode == "burst":
            # One warm connection, as an utterance would arrive on the open WebSocket.
            while line := readline(conn):
                left = int(line)
                while left > 0:
                    n = conn.send(CHUNK[:min(left, len(CHUNK))])
                    left -= n
                    sent += n
        elif mode == "paced":
            rate, secs = float(args[1]), float(args[2])
            size = max(12, int(rate / 30))
            t0 = time.monotonic()
            for i in range(int(secs * 30)):
                # Pace on the schedule, not on the last send: a stall must not lower the offered rate.
                delay = t0 + i / 30 - time.monotonic()
                if delay > 0:
                    time.sleep(delay)
                conn.sendall(struct.pack("!Id", size, time.time()) + bytes(size - 12))
                sent += size
        else:
            end = time.monotonic() + float(args[0])
            while time.monotonic() < end:
                sent += conn.send(CHUNK)
    except OSError as e:
        log({"id": test_id, "kind": "tcp_down_orin", "mode": mode, "error": str(e)})
    log({"id": test_id, "kind": "tcp_down_orin", "mode": mode, "args": args, "pace": PACE, "peer": peer, "bytes": sent,
         **tcp_info(conn)})


def up(conn, peer):
    test_id = readline(conn)
    got = 0
    t0 = time.monotonic()
    while True:
        b = conn.recv(65536)
        if not b:
            break
        got += len(b)
    secs = time.monotonic() - t0
    conn.sendall(f"{got}\n".encode())
    log({"id": test_id, "kind": "tcp_up_orin", "peer": peer, "bytes": got, "secs": round(secs, 3),
         **tcp_info(conn)})


def results(conn, peer):
    data = b""
    while True:
        b = conn.recv(65536)
        if not b:
            break
        data += b
    for line in data.decode().splitlines():
        if line.strip():
            log({"peer": peer, **json.loads(line)})


def serve_tcp(port, handler):
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    s.bind(("0.0.0.0", port))
    s.listen(4)
    while True:
        conn, addr = s.accept()

        def run(conn=conn, peer=addr[0]):
            with conn:
                try:
                    handler(conn, peer)
                except Exception as e:  # a broken test must not stop the server
                    print(f"{handler.__name__} {peer}: {e}", flush=True)

        threading.Thread(target=run, daemon=True).start()


def serve_echo(port):
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.bind(("0.0.0.0", port))
    while True:
        data, addr = s.recvfrom(2048)
        s.sendto(data, addr)


if __name__ == "__main__":
    os.makedirs(os.path.dirname(RESULTS), exist_ok=True)
    for port, h in ((7791, down), (7792, up), (7794, results)):
        threading.Thread(target=serve_tcp, args=(port, h), daemon=True).start()
    print(f"linktest server: tcp 7791 down, 7792 up, 7794 results; udp 7793 echo; -> {RESULTS}", flush=True)
    serve_echo(7793)