#!/usr/bin/env python3
"""UDP Stream black-box test: sends multiple packets from each stable client socket.

Start separately labeled mock backends and target an existing RustProxy UDP
listener. 'selftest' only validates this Python harness, NOT RustProxy.
By default, the server responds PONG to PING for an explicitly configured probe.
"""
from __future__ import annotations

import argparse
from collections import Counter
from concurrent.futures import ThreadPoolExecutor, as_completed
import secrets
import socket
import socketserver
import sys
import threading
import time

MAX_DATAGRAM = 1024


class UDPServer(socketserver.UDPServer):
    allow_reuse_address = True


class UDPBackend(socketserver.BaseRequestHandler):
    def handle(self) -> None:
        data, connection = self.request
        if data == b'PING':
            reply = b'PONG'
        elif data and len(data) <= MAX_DATAGRAM:
            reply = self.server.backend_id.encode('ascii') + b'|' + data
        else:
            return
        connection.sendto(reply, self.client_address)


def run_server(host: str, port: int, backend_id: str) -> None:
    with UDPServer((host, port), UDPBackend) as server:
        server.backend_id = backend_id
        print(f'UDP mock {backend_id} listening on {server.server_address}; health probe PING -> PONG', flush=True)
        try:
            server.serve_forever(poll_interval=0.1)
        except KeyboardInterrupt:
            pass


def check_one(index: int, host: str, port: int, messages: int, timeout: float, pause_ms: float, nonce: str) -> str:
    info = socket.getaddrinfo(host, port, type=socket.SOCK_DGRAM)[0]
    family, sock_type, proto, _, address = info
    backend = None
    # One socket = one stable client identity for the entire UDP Stream session.
    with socket.socket(family, sock_type, proto) as connection:
        connection.settimeout(timeout)
        for seq in range(messages):
            payload = f'rp-udp:{nonce}:{index}:{seq}'.encode('ascii')
            connection.sendto(payload, address)
            response, _ = connection.recvfrom(MAX_DATAGRAM + 128)
            try:
                observed, echoed = response.split(b'|', 1)
            except ValueError as exc:
                raise AssertionError(f'session {index}: missing mock backend response: {response!r}') from exc
            assert echoed == payload, f'session {index}: response mismatch {response!r}'
            assert observed, f'session {index}: empty backend ID'
            if backend is None:
                backend = observed
            assert observed == backend, f'session {index}: UDP backend changed for the same client socket'
            if pause_ms and seq != messages - 1:
                time.sleep(pause_ms / 1000)
    return backend.decode('ascii')


def run_client(host: str, port: int, sessions: int, messages: int, parallel: int, timeout: float,
               pause_ms: float, min_backends: int) -> Counter:
    nonce = secrets.token_hex(8)
    started = time.monotonic()
    counts: Counter[str] = Counter()
    with ThreadPoolExecutor(max_workers=parallel) as pool:
        futures = [pool.submit(check_one, i, host, port, messages, timeout, pause_ms, nonce)
                   for i in range(sessions)]
        for future in as_completed(futures):
            counts[future.result()] += 1
    assert len(counts) >= min_backends, f'observed {len(counts)} backend(s), required {min_backends}: {dict(counts)}'
    print(f'UDP PASS: {sessions} client sockets, {sessions * messages} datagrams, '
          f'{len(counts)} backend(s), {time.monotonic() - started:.2f}s; assignments={dict(sorted(counts.items()))}')
    return counts


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    sub = parser.add_subparsers(dest='mode', required=True)
    server = sub.add_parser('server', help='run local UDP echo and optional PING/PONG probe responder')
    server.add_argument('--host', default='127.0.0.1')
    server.add_argument('--port', type=int, required=True)
    server.add_argument('--backend-id', required=True, help='unique ASCII marker (e.g. node-a)')
    client = sub.add_parser('client', help='test forwarding via an already-running RustProxy UDP listener')
    client.add_argument('--host', default='127.0.0.1', help='RustProxy UDP listener host')
    client.add_argument('--port', type=int, required=True, help='RustProxy UDP listener port')
    client.add_argument('--sessions', type=int, default=24)
    client.add_argument('--messages', type=int, default=3)
    client.add_argument('--parallel', type=int, default=8)
    client.add_argument('--timeout', type=float, default=4.0)
    client.add_argument('--pause-ms', type=float, default=20.0, help='pause between datagrams on each stable client socket')
    client.add_argument('--min-backends', type=int, default=1, help='minimum distinct backend IDs observed')
    sub.add_parser('selftest', help='test mock peer + harness directly, WITHOUT RustProxy')
    args = parser.parse_args()
    if args.mode == 'server':
        if not args.backend_id.isascii() or not args.backend_id or '|' in args.backend_id or '\n' in args.backend_id:
            parser.error('--backend-id must be nonempty ASCII without | or newline')
        run_server(args.host, args.port, args.backend_id)
    elif args.mode == 'client':
        if min(args.port, args.sessions, args.messages, args.parallel, args.min_backends) < 1 or args.port > 65535 or args.timeout <= 0 or args.pause_ms < 0:
            parser.error('invalid port, count, timeout or pause delay')
        run_client(args.host, args.port, args.sessions, args.messages, args.parallel,
                   args.timeout, args.pause_ms, args.min_backends)
    else:
        with UDPServer(('127.0.0.1', 0), UDPBackend) as server:
            server.backend_id = 'local-mock'
            thread = threading.Thread(target=server.serve_forever, daemon=True)
            thread.start()
            try:
                run_client('127.0.0.1', server.server_address[1], 8, 3, 4, 2.0, 5, 1)
                print('SELFTEST PASS: tested mock UDP peer, not RustProxy')
            finally:
                server.shutdown()
                thread.join(timeout=2)


if __name__ == '__main__':
    try:
        main()
    except (OSError, AssertionError, UnicodeError) as exc:
        print(f'UDP FAIL: {exc}', file=sys.stderr)
        sys.exit(1)