"""Count socket syscalls for identical tiny requests; not a throughput benchmark."""
import argparse
import hashlib
import http.client
import json
import os
from pathlib import Path
import signal
import socket
import subprocess
import sys
import time

root = Path(__file__).resolve().parents[1]
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--server', choices=['skirr', 'peregrine', 'weft'], required=True)
parser.add_argument('--baseline', type=Path)
parser.add_argument('--app', choices=['raw', 'fastapi'], default='raw')
parser.add_argument('--output', type=Path, required=True)
a = parser.parse_args()
source = a.baseline / 'python' if a.baseline else root / 'python'
out = a.output.resolve()
out.mkdir(parents=True, exist_ok=True)
with socket.socket() as sock:
    sock.bind(('127.0.0.1', 0))
    port = sock.getsockname()[1]
env = dict(os.environ, PYTHONPATH=str(source) + os.pathsep + str(root / 'bench'))
command = ['strace', '-f', '-qq', '-yy', '-e', 'trace=network,read,write,writev,epoll_ctl,epoll_wait',
           '-o', str(out / 'syscalls.txt'), sys.executable, str(root / 'bench/server.py'),
           a.server, '--app', a.app, '--port', str(port), '--workers', '1', '--cpus', '0']
(out / 'command.json').write_text(json.dumps(command, indent=2))
path, expected = ('/json', b'{"hello":"world"}') if a.app == 'fastapi' else ('/tiny', b'Hello world!')
process = subprocess.Popen(command, env=env, stdout=subprocess.DEVNULL,
                           stderr=subprocess.DEVNULL, start_new_session=True)
try:
    for _ in range(100):
        client = http.client.HTTPConnection('127.0.0.1', port, timeout=.5)
        try:
            client.request('GET', path)
            response = client.getresponse()
            assert response.status == 200 and response.read() == expected
            break
        except OSError:
            time.sleep(.05)
        finally:
            client.close()
    else:
        raise RuntimeError('Server did not start')
    client = http.client.HTTPConnection('127.0.0.1', port, timeout=5)
    try:
        for _ in range(1000):
            client.request('GET', path)
            response = client.getresponse()
            assert response.status == 200 and response.read() == expected
    finally:
        client.close()
finally:
    os.killpg(process.pid, signal.SIGINT)
    try:
        process.wait(timeout=5)
    except subprocess.TimeoutExpired:
        os.killpg(process.pid, signal.SIGKILL)
        process.wait()
lines = (out / 'syscalls.txt').read_text().splitlines()
reads = [line for line in lines if '<TCP:' in line and
         ('read(' in line or 'recvfrom(' in line)]
writes = [line for line in lines if '<TCP:' in line and
          any(name in line for name in ('write(', 'sendto(', 'writev('))]
summary = dict(server=a.server, app=a.app, requests=1001, socket_reads=len(reads),
               epoll_out_modifications=sum("epoll_ctl(" in line and "EPOLL_CTL_MOD" in line and "EPOLLOUT" in line for line in lines),
               socket_error_queries=sum("getsockopt(" in line and "SO_ERROR" in line for line in lines),
               nodelay_pushes=sum("setsockopt(" in line and "TCP_NODELAY" in line for line in lines),
               source_sha256={p.name: hashlib.sha256(p.read_bytes()).hexdigest() for p in (a.baseline or root).joinpath("src").glob("*.rs")} if a.server == "skirr" else {},
               binary_sha256=hashlib.sha256((source/"skirr/_native.so").read_bytes()).hexdigest() if a.server == "skirr" else None,
               socket_read_EAGAIN=sum('EAGAIN' in line for line in reads),
               socket_writes=len(writes), diagnostic_only=True,
               includes_startup_and_shutdown=True)
(out / 'summary.json').write_text(json.dumps(summary, indent=2))
print(summary)