#!/usr/bin/env python3
"""Small dependency-free UDP path test with sequence, loss, and RTT reporting.

Run a bounded echo receiver on the remote host:
    python3 scripts/netdiag/udp_path.py server --bind 0.0.0.0 --port 443 --duration 45

Then run the client from the source host:
    python3 scripts/netdiag/udp_path.py client --target 203.0.113.10 --port 443 \
        --count 251 --rate 31 --size 64 --timeout 3 --socket-mode sendto

The server echoes each valid datagram unchanged. The client uses its sequence and
monotonic send timestamp to report lost replies and round-trip latency. It uses
no TCP control channel, so TCP service occupancy on the same port is irrelevant.
"""

from __future__ import annotations

import argparse
import math
import secrets
import select
import socket
import struct
import sys
import time
from collections import Counter


HEADER = struct.Struct("!8sIQ")  # random run id, sequence, monotonic ns
MAX_DATAGRAM = 65_507


def percentile(values: list[float], p: float) -> float | None:
    if not values:
        return None
    ordered = sorted(values)
    index = max(0, math.ceil(p * len(ordered)) - 1)
    return ordered[index]


def format_ms(value: float | None) -> str:
    return "n/a" if value is None else f"{value:.3f} ms"


def is_probe_packet(data: bytes) -> bool:
    return len(data) >= HEADER.size


def run_server(args: argparse.Namespace) -> int:
    sock = socket.socket(socket.AF_INET6 if ":" in args.bind else socket.AF_INET, socket.SOCK_DGRAM)
    if sock.family == socket.AF_INET6:
        sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0)
    try:
        sock.bind((args.bind, args.port))
    except OSError as exc:
        print(f"bind failed on {args.bind}:{args.port}: {exc}", file=sys.stderr)
        return 2

    print(f"listening udp://{args.bind}:{args.port} for {args.duration:g}s", flush=True)
    deadline = time.monotonic() + args.duration if args.duration > 0 else None
    received = echoed = malformed = duplicates = 0
    sequences: set[tuple[str, bytes, int]] = set()
    sources: Counter[str] = Counter()
    next_report = time.monotonic() + 1
    sock.settimeout(0.25)
    try:
        while deadline is None or time.monotonic() < deadline:
            try:
                data, address = sock.recvfrom(MAX_DATAGRAM)
            except socket.timeout:
                continue
            received += 1
            source = f"{address[0]}:{address[1]}"
            sources[source] += 1
            if not is_probe_packet(data):
                malformed += 1
                continue
            run_id, sequence, _ = HEADER.unpack_from(data)
            key = (source, run_id, sequence)
            if key in sequences:
                duplicates += 1
            else:
                sequences.add(key)
            try:
                sock.sendto(data, address)
                echoed += 1
            except OSError as exc:
                print(f"reply failed to {source}: {exc}", file=sys.stderr, flush=True)
            if time.monotonic() >= next_report:
                print(f"server progress received={received} unique={len(sequences)} echoed={echoed}", flush=True)
                next_report += 1
    except KeyboardInterrupt:
        pass
    finally:
        sock.close()

    source_summary = ", ".join(f"{source}={count}" for source, count in sources.most_common()) or "none"
    print(
        f"server received={received} unique={len(sequences)} echoed={echoed} "
        f"malformed={malformed} duplicate_seq={duplicates} sources=[{source_summary}]",
        flush=True,
    )
    return 0


def receive_replies(
    sock: socket.socket,
    pending: dict[int, int],
    rtts_ms: list[float],
    duplicates: list[int],
    wrong_source: list[int],
    receive_errors: list[str],
    run_id: bytes,
    target: tuple[str, int],
    wait_seconds: float,
) -> None:
    deadline = time.monotonic() + wait_seconds
    while time.monotonic() < deadline:
        remaining = deadline - time.monotonic()
        readable, _, _ = select.select([sock], [], [], max(0.0, min(remaining, 0.05)))
        if not readable:
            continue
        try:
            data, address = sock.recvfrom(MAX_DATAGRAM)
        except OSError as exc:
            if not receive_errors:
                receive_errors.append(f"{type(exc).__name__}: {exc}")
            continue
        if address[1] != target[1] or address[0] not in (target[0], "::ffff:" + target[0]):
            wrong_source[0] += 1
            continue
        if not is_probe_packet(data):
            continue
        reply_run_id, sequence, sent_ns = HEADER.unpack_from(data)
        if reply_run_id != run_id:
            wrong_source[0] += 1
            continue
        if sequence not in pending:
            duplicates[0] += 1
            continue
        pending.pop(sequence)
        rtts_ms.append((time.monotonic_ns() - sent_ns) / 1_000_000)


def run_client(args: argparse.Namespace) -> int:
    if args.size < HEADER.size or args.size > MAX_DATAGRAM:
        print(f"--size must be between {HEADER.size} and {MAX_DATAGRAM}", file=sys.stderr)
        return 2
    try:
        addresses = socket.getaddrinfo(args.target, args.port, type=socket.SOCK_DGRAM)
    except OSError as exc:
        print(f"resolve failed for {args.target}: {exc}", file=sys.stderr)
        return 2
    family, socktype, proto, _, target = addresses[0]
    sock = socket.socket(family, socktype, proto)
    pending: dict[int, int] = {}
    rtts_ms: list[float] = []
    duplicates = [0]
    wrong_source = [0]
    receive_errors: list[str] = []
    send_errors = 0
    run_id = secrets.token_bytes(8)
    local_endpoint = None
    payload_padding = bytes(args.size - HEADER.size)

    try:
        if args.socket_mode == "connected":
            sock.connect(target)
        start = time.monotonic()
        for sequence in range(args.count):
            send_at = start + sequence / args.rate
            while True:
                remaining = send_at - time.monotonic()
                if remaining <= 0:
                    break
                receive_replies(sock, pending, rtts_ms, duplicates, wrong_source, receive_errors, run_id, target, max(0.0, min(remaining, 0.05)))

            sent_ns = time.monotonic_ns()
            packet = HEADER.pack(run_id, sequence, sent_ns) + payload_padding
            try:
                if args.socket_mode == "connected":
                    sock.send(packet)
                else:
                    sock.sendto(packet, target)
                if local_endpoint is None:
                    local_endpoint = sock.getsockname()
                pending[sequence] = sent_ns
            except OSError as exc:
                send_errors += 1
                print(f"send failed seq={sequence}: {exc}", file=sys.stderr, flush=True)

        receive_replies(sock, pending, rtts_ms, duplicates, wrong_source, receive_errors, run_id, target, args.timeout)
    finally:
        sock.close()

    sent = args.count - send_errors
    received = len(rtts_ms)
    loss = 0.0 if sent == 0 else 100.0 * (sent - received) / sent
    elapsed = args.count / args.rate
    print(f"target={args.target}:{args.port} local={local_endpoint} socket_mode={args.socket_mode} size={args.size} rate={args.rate:g} pps duration≈{elapsed:.2f}s run_id={run_id.hex()}")
    print(f"sent={sent} replies={received} timeout_loss={loss:.2f}% send_errors={send_errors} duplicates={duplicates[0]} wrong_source={wrong_source[0]} receive_errors={receive_errors}")
    print(f"rtt p50={format_ms(percentile(rtts_ms, 0.50))} p95={format_ms(percentile(rtts_ms, 0.95))} max={format_ms(max(rtts_ms) if rtts_ms else None)}")
    return 0 if received > 0 else 1


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    modes = parser.add_subparsers(dest="mode", required=True)
    server = modes.add_parser("server", help="run a bounded or continuous UDP echo receiver")
    server.add_argument("--bind", default="0.0.0.0")
    server.add_argument("--port", type=int, required=True)
    server.add_argument("--duration", type=float, default=45, help="seconds; 0 runs until Ctrl-C")
    server.set_defaults(run=run_server)
    client = modes.add_parser("client", help="send sequenced UDP probes and measure replies")
    client.add_argument("--target", required=True)
    client.add_argument("--port", type=int, required=True)
    client.add_argument("--count", type=int, default=251)
    client.add_argument("--rate", type=float, default=31, help="datagrams per second")
    client.add_argument("--size", type=int, default=64, help="UDP payload size in bytes")
    client.add_argument("--timeout", type=float, default=3, help="reply drain timeout in seconds")
    client.add_argument(
        "--socket-mode",
        choices=("sendto", "connected"),
        default="sendto",
        help="send with sendto() or a connected UDP socket; compare these when diagnosing VM networking",
    )
    client.set_defaults(run=run_client)
    return parser


if __name__ == "__main__":
    arguments = build_parser().parse_args()
    if getattr(arguments, "count", 1) < 1 or getattr(arguments, "rate", 1) <= 0:
        raise SystemExit("--count must be positive and --rate must be greater than zero")
    raise SystemExit(arguments.run(arguments))