#!/usr/bin/env python3
"""Repeatable WebSocket game traffic. Requires Python 3.12, websockets 17.1, psutil.

uv run --python 3.12 --with websockets==17.1 --with psutil \
  python bench/game_traffic.py run --binary ./superweb --output build/game.json
A separate backend process handles RFC 6455; the benchmark checks echoed bytes.
"""
import argparse
import asyncio
import collections
import contextlib
import hashlib
import json
import os
import pathlib
import platform
import signal
import socket
import ssl
import struct
import subprocess
import sys
import tempfile
import time

import psutil
import websockets
from websockets.asyncio.client import connect
from websockets.asyncio.server import serve


def percentile(values, p):
    if not values:
        return None
    values = sorted(values)
    return values[min(len(values) - 1, int((len(values) - 1) * p))]


def distribution(values):
    return {name: percentile(values, p) for name, p in [('p50', .5), ('p95', .95), ('p99', .99), ('max', 1)]}


async def backend(port, host='127.0.0.1'):
    rooms = collections.defaultdict(set)
    async def handler(ws):
        path = ws.request.path
        room = path if path.startswith('/room/') else None
        if room:
            rooms[room].add(ws)
        try:
            if path == '/flood':
                while True:
                    await ws.send(b'x' * 65536)
                    await asyncio.sleep(0)
            async for message in ws:
                if room:
                    # Same deterministic broadcast backend for both binaries.
                    # Each recipient has its own write task so a slow socket
                    # cannot serialize the other clients in the room.
                    await asyncio.gather(*(peer.send(message) for peer in tuple(rooms[room])), return_exceptions=True)
                else:
                    await ws.send(message)
        except websockets.ConnectionClosed:
            pass
        finally:
            if room:
                rooms[room].discard(ws)
    async with serve(handler, host, port, compression=None, ping_interval=5, ping_timeout=10,
                     close_timeout=1, max_size=2**20, max_queue=16, write_limit=32768):
        print('ready', flush=True)
        await asyncio.Future()


def unused_port():
    with socket.socket() as sock:
        sock.bind(('127.0.0.1', 0))
        return sock.getsockname()[1]


def context(cafile=None):
    ctx = ssl.create_default_context(cafile=cafile)
    ctx.set_alpn_protocols(['http/1.1'])
    return ctx


async def exercise(url, count, size, hz, seconds, kind, cafile=None, asset_url=None):
    ctx = context(cafile) if url.startswith('wss:') else None
    errors = collections.Counter()
    sockets = []
    handshakes = []
    metrics = {'attempted_connections': count, 'requested_hz_per_player': hz,
               'payload_bytes': size, 'seconds': seconds, 'kind': kind}
    async def dial(i):
        start = time.perf_counter()
        target = url + (f'/room/{i // 16}' if kind == 'broadcast' else '/flood' if kind == 'slow' and i % 8 == 0 else '/echo')
        try:
            ws = await connect(target, ssl=ctx, proxy=None, compression=None, open_timeout=3,
                               close_timeout=.3, ping_interval=None, max_queue=16, max_size=2**20)
            handshakes.append((time.perf_counter() - start) * 1000)
            return i, ws
        except Exception as exc:
            errors['connect:' + type(exc).__name__] += 1
            return i, None
    sockets = [(i, ws) for i, ws in await asyncio.gather(*(dial(i) for i in range(count))) if ws]
    metrics['established_connections'] = len(sockets)
    if kind == 'reconnect':
        completed = 0
        for cycle in range(5):
            if cycle:
                sockets = [(i, ws) for i, ws in await asyncio.gather(*(dial(i) for i in range(count))) if ws]
                metrics['established_connections'] += len(sockets)
            async def exchange(i, ws):
                nonlocal completed
                try:
                    payload = struct.pack('!I', i) + b'r' * (size - 4)
                    await ws.send(payload)
                    if await asyncio.wait_for(ws.recv(), 2) != payload:
                        raise ValueError('reconnect payload')
                    completed += 1
                except Exception as exc:
                    errors['traffic:' + type(exc).__name__] += 1
                finally:
                    await ws.close()
            await asyncio.gather(*(exchange(i, ws) for i, ws in sockets))
        metrics.update(attempted_connections=count * 5, sent=completed, received=completed, verified=completed,
                       delivered_per_second=None, errors=dict(errors), handshake_ms=distribution(handshakes),
                       latency_ms=distribution([]), client_schedule_lag_ms=distribution([]), asset_latency_ms=distribution([]), asset_requests=0)
        return metrics
    latencies, lags, asset_latencies = [], [], []
    sent = received = verified = 0
    start = time.perf_counter() + .1
    stop = start + seconds
    async def player(i, ws):
        nonlocal sent, received, verified
        pending = {}
        async def receive():
            nonlocal received, verified
            try:
                while True:
                    data = await ws.recv()
                    now = time.perf_counter()
                    if not isinstance(data, bytes) or len(data) != size:
                        errors['invalid_payload'] += 1
                        continue
                    owner, seq, stamped = struct.unpack('!IId', data[:16])
                    expected = bytes([owner % 251]) * (size - 16)
                    if data[16:] != expected:
                        errors['invalid_payload'] += 1
                        continue
                    received += 1
                    verified += 1
                    latencies.append((now - stamped) * 1000)
                    if owner == i:
                        pending.pop(seq, None)
            except websockets.ConnectionClosed as exc:
                if time.perf_counter() < stop:
                    errors['disconnect:' + str(exc.rcvd.code if exc.rcvd else 'abnormal')] += 1
        slow_reader = kind == 'slow' and i % 8 == 0
        task = asyncio.create_task(asyncio.sleep(seconds + 1) if slow_reader else receive())
        try:
            if kind == 'idle':
                await asyncio.sleep(seconds)
                pong = await ws.ping()
                await asyncio.wait_for(pong, 1)
            elif kind == 'slow' and i % 8 == 0:
                task.cancel()
                # Fill this player's outbound stream while it stops reading.
                ws.transport.pause_reading()
                await asyncio.sleep(seconds)
                ws.transport.resume_reading()
            else:
                for seq in range(int(seconds * hz)):
                    due = start + seq / hz
                    await asyncio.sleep(max(0, due - time.perf_counter()))
                    now = time.perf_counter()
                    if now >= stop:
                        break
                    lags.append((now - due) * 1000)
                    if len(pending) >= 128:
                        errors['inflight_limit'] += 1
                        continue
                    data = struct.pack('!IId', i, seq, now) + bytes([i % 251]) * (size - 16)
                    pending[seq] = now
                    await asyncio.wait_for(ws.send(data), 1)
                    sent += 1
                await asyncio.sleep(max(0, stop - time.perf_counter()) + .3)
                if pending:
                    errors['unanswered_messages'] += len(pending)
        except Exception as exc:
            errors['traffic:' + type(exc).__name__] += 1
        finally:
            task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await task
    async def assets():
        if not asset_url:
            return
        from urllib.parse import urlsplit
        u = urlsplit(asset_url)
        while time.perf_counter() < stop:
            at = time.perf_counter()
            try:
                reader, writer = await asyncio.wait_for(asyncio.open_connection(
                    u.hostname, u.port or (443 if u.scheme == 'https' else 80),
                    ssl=context(cafile) if u.scheme == 'https' else None), 1)
                try:
                    writer.write(f'GET {u.path or "/"} HTTP/1.1\r\nHost: {u.hostname}\r\nConnection: close\r\n\r\n'.encode())
                    await writer.drain()
                    data = await asyncio.wait_for(reader.read(), 2)
                    if not data.startswith(b'HTTP/1.1 200 '):
                        raise ValueError('asset status')
                    head, sep, body = data.partition(b'\r\n\r\n')
                    if not sep or body != bytes(range(256)) * 4096:
                        raise ValueError('asset bytes')
                    asset_latencies.append((time.perf_counter() - at) * 1000)
                finally:
                    writer.close()
                    with contextlib.suppress(Exception):
                        await writer.wait_closed()
            except Exception as exc:
                errors['asset:' + type(exc).__name__] += 1
            await asyncio.sleep(.05)
    await asyncio.gather(*(player(i, ws) for i, ws in sockets), assets())
    await asyncio.gather(*(ws.close() for _, ws in sockets), return_exceptions=True)
    metrics.update(sent=sent, received=received, verified=verified,
                   delivered_per_second=received / seconds, errors=dict(errors),
                   latency_ms=distribution(latencies), handshake_ms=distribution(handshakes),
                   client_schedule_lag_ms=distribution(lags), asset_latency_ms=distribution(asset_latencies),
                   asset_requests=len(asset_latencies))
    return metrics


def stop_process(p):
    if p and p.poll() is None:
        p.terminate()
        try:
            p.wait(timeout=5)
        except subprocess.TimeoutExpired:
            p.kill()
            p.wait()


async def run(args):
    output = pathlib.Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    result = {'started': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime()),
              'environment': {'platform': platform.platform(), 'cpu': platform.processor(),
                              'cpus': os.cpu_count(), 'websockets': websockets.__version__},
              'method': 'Open-loop messages; concurrent 3s connection deadlines; latencies include only delivered messages. Failures are retained. No compression. TLS verified with generated local CA.',
              'samples': []}
    binary = str(pathlib.Path(args.binary).resolve())
    result['binary_sha256'] = hashlib.sha256(pathlib.Path(binary).read_bytes()).hexdigest()
    cases = [('echo-single', 1, 64, 60, 'echo'), ('echo-64', 64, 64, 60, 'echo'),
             ('echo-1024', 64, 1024, 20, 'echo'), ('broadcast', 64, 64, 20, 'broadcast'),
             ('mixed-assets', 64, 256, 60, 'echo'), ('idle', 16, 64, 20, 'idle'),
             ('slow-readers', 64, 256, 60, 'slow'), ('reconnect', 64, 64, 0, 'reconnect')]
    if args.cases:
        cases = [c for c in cases if c[0] in args.cases.split(',')]
    with tempfile.TemporaryDirectory(prefix='superweb-game-') as tmp:
        tmp = pathlib.Path(tmp)
        (tmp / 'asset.bin').write_bytes(bytes(range(256)) * 4096)
        cert, key = tmp / 'cert.pem', tmp / 'key.pem'
        subprocess.run(['openssl', 'req', '-x509', '-newkey', 'rsa:2048', '-nodes', '-days', '1',
                        '-keyout', str(key), '-out', str(cert), '-subj', '/CN=localhost',
                        '-addext', 'subjectAltName=DNS:localhost,IP:127.0.0.1'], check=True,
                       stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        for round_no in range(1, args.rounds + 1):
            for name, count, size, hz, kind in cases:
                backend_port, port, tls_port = unused_port(), unused_port(), unused_port()
                back = server = None
                log = output.parent / f'{output.stem}-{round_no}-{name}.log'
                with log.open('w') as server_log:
                    try:
                        back = subprocess.Popen([sys.executable, __file__, 'backend', '--port', str(backend_port)],
                                                stdout=subprocess.PIPE, stderr=server_log, text=True)
                        assert await asyncio.to_thread(back.stdout.readline) == 'ready\n'
                        command = [binary, '--root', str(tmp), '--bind', '127.0.0.1', '--port', str(port),
                                   '--workers', '2', '--tls-port', str(tls_port), '--tls-cert', str(cert),
                                   '--tls-key', str(key), '--proxy', f'/game=127.0.0.1:{backend_port}']
                        server = subprocess.Popen(command, stdout=server_log, stderr=server_log)
                        await asyncio.sleep(.25)
                        if server.poll() is not None:
                            raise RuntimeError(log.read_text())
                        (output.parent / 'active-server.pid').write_text(str(server.pid))
                        monitored = {'server': psutil.Process(server.pid), 'backend': psutil.Process(back.pid), 'client': psutil.Process()}
                        snapshots = []
                        async def monitor():
                            while True:
                                row = {'elapsed': time.perf_counter()}
                                for label, process in monitored.items():
                                    try:
                                        m = process.memory_full_info()
                                    except psutil.AccessDenied:
                                        m = process.memory_info()
                                    cpu = process.cpu_times()
                                    row[label] = {'rss': m.rss, 'pss': getattr(m, 'pss', None),
                                                  'cpu_seconds': cpu.user + cpu.system}
                                snapshots.append(row)
                                await asyncio.sleep(.2)
                        task = asyncio.create_task(monitor())
                        row = await exercise(f'wss://localhost:{tls_port}/game', count, size, hz,
                                             max(12, args.seconds) if kind == 'idle' else args.seconds,
                                             kind, str(cert), f'https://localhost:{tls_port}/asset.bin' if name == 'mixed-assets' else None)
                        task.cancel()
                        with contextlib.suppress(asyncio.CancelledError):
                            await task
                        row.update(case=name, round=round_no, resources=snapshots)
                        result['samples'].append(row)
                        output.write_text(json.dumps(result, indent=2) + '\n')
                        print(json.dumps({k:v for k,v in row.items() if k != 'resources'}), flush=True)
                    finally:
                        stop_process(server)
                        stop_process(back)
    result['finished'] = time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())
    output.write_text(json.dumps(result, indent=2) + '\n')


async def load(args):
    cases = [('echo-single', 1, 64, 60, 'echo'), ('echo-64', 64, 64, 60, 'echo'),
             ('echo-1024', 64, 1024, 20, 'echo'), ('broadcast', 64, 64, 20, 'broadcast'),
             ('mixed-assets', 64, 256, 60, 'echo'), ('idle', 16, 64, 20, 'idle'),
             ('slow-readers', 64, 256, 60, 'slow'), ('reconnect', 64, 64, 0, 'reconnect')]
    if args.connections:
        cases = [('capacity', args.connections, 256, 20, 'echo')]
    if args.cases:
        cases = [c for c in cases if c[0] in args.cases.split(',')]
    result = {'started': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime()), 'samples': [],
              'environment': {'platform': platform.platform(), 'cpus': os.cpu_count(), 'websockets': websockets.__version__},
              'method': 'Verified TLS; uncompressed RFC 6455 messages; concurrent 3s connection deadlines. Message latency includes only delivered messages; connection failures and unanswered messages are retained. Room size 16.'}
    process = psutil.Process()
    for number in range(1, args.rounds + 1):
        for name, count, size, hz, kind in cases:
            before = process.cpu_times(); start = time.perf_counter()
            started = time.time()
            row = await exercise(args.url, count, size, hz, max(12, args.seconds) if kind == 'idle' else args.seconds,
                                 kind, args.cafile, args.asset_url if name == 'mixed-assets' else None)
            after = process.cpu_times()
            row.update(case=name, round=number, started_epoch=started, finished_epoch=time.time(),
                       client_cpu_cores=(after.user + after.system - before.user - before.system)/(time.perf_counter()-start),
                       client_rss_mib=process.memory_info().rss/1024**2)
            result['samples'].append(row)
            pathlib.Path(args.output).write_text(json.dumps(result, indent=2)+'\n')
            print(json.dumps(row), flush=True)
    result['finished'] = time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())
    pathlib.Path(args.output).write_text(json.dumps(result, indent=2)+'\n')


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    sub = parser.add_subparsers(dest='action', required=True)
    p = sub.add_parser('backend')
    p.add_argument('--port', type=int, required=True)
    p.add_argument('--host', default='127.0.0.1')
    p = sub.add_parser('run')
    p.add_argument('--binary', required=True)
    p.add_argument('--output', required=True)
    p.add_argument('--seconds', type=int, default=8)
    p.add_argument('--rounds', type=int, default=3)
    p.add_argument('--cases')
    p = sub.add_parser('load')
    p.add_argument('--url', required=True)
    p.add_argument('--output', required=True)
    p.add_argument('--cafile')
    p.add_argument('--asset-url')
    p.add_argument('--seconds', type=int, default=8)
    p.add_argument('--rounds', type=int, default=3)
    p.add_argument('--connections', type=int)
    p.add_argument('--cases')
    args = parser.parse_args()
    asyncio.run(backend(args.port, args.host) if args.action == 'backend' else load(args) if args.action == 'load' else run(args))


if __name__ == '__main__':
    main()