#!/usr/bin/env python3
"""Real RFC 6455 handshakes/frames, WSS, backpressure and worker fairness (stdlib)."""
import asyncio
import base64
import concurrent.futures
import contextlib
import hashlib
import os
import pathlib
import re
import signal
import socket
import ssl
import struct
import subprocess
import tempfile
import threading
import time
import unittest

ROOT = pathlib.Path(__file__).resolve().parents[1]
BINARY = pathlib.Path(os.environ.get('SUPERWEB', ROOT / 'superweb')).resolve()


def frame(data, opcode=2, fin=True, masked=True):
    n = len(data)
    head = bytes([(128 if fin else 0) | opcode, (128 if masked else 0) | (n if n < 126 else 126 if n < 65536 else 127)])
    if n >= 126:
        head += struct.pack('!H' if n < 65536 else '!Q', n)
    if masked:
        key = b'game'
        return head + key + bytes(v ^ key[i % 4] for i, v in enumerate(data))
    return head + data


def exact(sock, n):
    out = b''
    while len(out) < n:
        chunk = sock.recv(n - len(out))
        if not chunk:
            raise EOFError()
        out += chunk
    return out


def recv_frame(sock):
    a, b = exact(sock, 2)
    n = b & 127
    if n == 126:
        n = struct.unpack('!H', exact(sock, 2))[0]
    elif n == 127:
        n = struct.unpack('!Q', exact(sock, 8))[0]
    assert not b & 128
    return a, exact(sock, n)


class Backend:
    def __init__(self):
        self.loop = asyncio.new_event_loop()
        self.ready = threading.Event()
        self.thread = threading.Thread(target=self.main, daemon=True)
        self.thread.start()
        assert self.ready.wait(5)

    def main(self):
        asyncio.set_event_loop(self.loop)
        self.server = self.loop.run_until_complete(asyncio.start_server(self.handle, '127.0.0.1', 0))
        self.port = self.server.sockets[0].getsockname()[1]
        self.ready.set()
        self.loop.run_forever()
        self.server.close()
        self.loop.run_until_complete(self.server.wait_closed())
        tasks = asyncio.all_tasks(self.loop)
        for task in tasks:
            task.cancel()
        self.loop.run_until_complete(asyncio.gather(*tasks, return_exceptions=True))
        self.loop.close()

    async def handle(self, reader, writer):
        try:
            header = await reader.readuntil(b'\r\n\r\n')
            path = header.split(b' ')[1]
            if path == b'/stall':
                await asyncio.sleep(3)
                return
            key = re.search(rb'(?im)^sec-websocket-key:\s*([^\r]+)', header)[1]
            accept = base64.b64encode(hashlib.sha1(key + b'258EAFA5-E914-47DA-95CA-C5AB0DC85B11').digest())
            response = b'HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: ' + accept + b'\r\n\r\n'
            if path == b'/greeting':
                response += frame(b'greeting', masked=False)
            writer.write(response)
            await writer.drain()
            if path == b'/flood':
                for _ in range(1024):
                    writer.write(frame(b'x' * 65536, masked=False))
                    await writer.drain()
                return
            while True:
                a, b = await reader.readexactly(2)
                n = b & 127
                if n == 126:
                    n = struct.unpack('!H', await reader.readexactly(2))[0]
                elif n == 127:
                    n = struct.unpack('!Q', await reader.readexactly(8))[0]
                if n > 2**21 or not b & 128:
                    return
                mask = await reader.readexactly(4)
                raw = await reader.readexactly(n)
                payload = bytes(v ^ mask[i % 4] for i, v in enumerate(raw))
                op = a & 15
                writer.write(frame(payload, 10 if op == 9 else op, bool(a & 128), False))
                await writer.drain()
                if op == 8:
                    return
        except (asyncio.IncompleteReadError, ConnectionError, OSError):
            pass
        finally:
            writer.close()
            with contextlib.suppress(ConnectionError, OSError):
                await writer.wait_closed()

    def close(self):
        self.loop.call_soon_threadsafe(self.loop.stop)
        self.thread.join(5)


class WebSocketTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.temp = tempfile.TemporaryDirectory(prefix='sw-ws-test-')
        cls.root = pathlib.Path(cls.temp.name)
        (cls.root / 'hello.txt').write_text('game assets work\n')
        cls.key, cls.cert = cls.root / 'key.pem', cls.root / 'cert.pem'
        subprocess.run(['openssl', 'req', '-x509', '-newkey', 'rsa:2048', '-nodes', '-days', '1', '-keyout', str(cls.key), '-out', str(cls.cert), '-subj', '/CN=localhost', '-addext', 'subjectAltName=DNS:localhost,IP:127.0.0.1'], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        cls.backend = Backend()
        cls.tls = '--tls-port' in subprocess.check_output([str(BINARY), '--help'], text=True) and 'x64' not in BINARY.name
        cls.ctx = ssl.create_default_context(cafile=str(cls.cert))
        cls.ctx.set_alpn_protocols(['http/1.1'])

    @classmethod
    def tearDownClass(cls):
        cls.backend.close()
        cls.temp.cleanup()

    def setUp(self):
        self.sockets = []
        self.processes = []

    def tearDown(self):
        for s in self.sockets:
            s.close()
        for p in self.processes:
            if p.poll() is None:
                p.terminate()
                try:
                    p.wait(timeout=4)
                except subprocess.TimeoutExpired:
                    p.kill(); p.wait()
            if p.stdout:
                p.stdout.close()
            if p.stderr:
                errors = p.stderr.read()
                self.assertNotIn('ERROR: AddressSanitizer', errors)
                self.assertNotIn('runtime error:', errors)
                p.stderr.close()

    def start(self, extra=(), tls=False):
        command = [str(BINARY), '--root', str(self.root), '--bind', '127.0.0.1', '--port', '0', '--workers', '1', '--proxy', f'/game=localhost:{self.backend.port}', *extra]
        if tls:
            command += ['--tls-port', '0', '--tls-cert', str(self.cert), '--tls-key', str(self.key)]
        p = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        self.processes.append(p)
        line = p.stdout.readline()
        self.assertIn('superweb listening', line)
        return p, int(re.search(r'https://[^:]+:(\d+)' if tls else r'http://[^:]+:(\d+)', line)[1])

    def connect(self, port, path='/echo', tls=False, early=b'', receive_buffer=0):
        s = socket.socket()
        s.settimeout(3)
        if receive_buffer:
            s.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, receive_buffer)
        s.connect(('127.0.0.1', port))
        if tls:
            s = self.ctx.wrap_socket(s, server_hostname='localhost')
        self.sockets.append(s)
        key = base64.b64encode(os.urandom(16))
        s.sendall(b'GET /game' + path.encode() + b' HTTP/1.1\r\nHost: localhost\r\nUpgrade: websocket\r\nConnection: keep-alive, Upgrade\r\nSec-WebSocket-Version: 13\r\nSec-WebSocket-Key: ' + key + b'\r\n\r\n' + early)
        header = b''
        while not header.endswith(b'\r\n\r\n'):
            header += exact(s, 1)
        expected = base64.b64encode(hashlib.sha1(key + b'258EAFA5-E914-47DA-95CA-C5AB0DC85B11').digest())
        self.assertTrue(header.startswith(b'HTTP/1.1 101'))
        self.assertIn(b'Sec-WebSocket-Accept: ' + expected, header)
        return s

    def asset(self, port):
        with socket.create_connection(('127.0.0.1', port), 2) as sock:
            sock.settimeout(2)
            sock.sendall(b'GET /hello.txt HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n')
            out = b''
            while chunk := sock.recv(4096):
                out += chunk
            self.assertTrue(out.endswith(b'game assets work\n'))

    def test_many_sessions_frames_and_assets(self):
        _, port = self.start()
        with concurrent.futures.ThreadPoolExecutor(max_workers=16) as pool:
            sessions = list(pool.map(lambda _: self.connect(port), range(64)))
        for index, sock in enumerate(sessions):
            sock.sendall(frame(struct.pack('!I', index)))
        self.asset(port)
        for index, sock in enumerate(sessions):
            self.assertEqual(recv_frame(sock), (130, struct.pack('!I', index)))
        sock = sessions[0]
        sock.sendall(frame(b'first', fin=False) + frame(b'last', opcode=0) + frame(b'ping', opcode=9))
        self.assertEqual(recv_frame(sock), (2, b'first'))
        self.assertEqual(recv_frame(sock), (128, b'last'))
        self.assertEqual(recv_frame(sock), (138, b'ping'))
        sock.sendall(frame(struct.pack('!H', 1000), opcode=8))
        self.assertEqual(recv_frame(sock), (136, struct.pack('!H', 1000)))
        self.assertEqual(sock.recv(1), b'')

    def test_tls_large_frames_and_coalesced_upgrade(self):
        if not self.tls:
            self.skipTest('TLS unavailable in Intel macOS build')
        _, port = self.start(['--ws-buffer', '4096'], tls=True)
        sock = self.connect(port, '/greeting', tls=True, early=frame(b'early'))
        self.assertEqual(recv_frame(sock), (130, b'greeting'))
        self.assertEqual(recv_frame(sock), (130, b'early'))
        body = os.urandom(512 * 1024)
        sock.sendall(frame(body))
        self.assertEqual(recv_frame(sock), (130, body))
        sock.sendall(frame(b'after'))
        self.assertEqual(recv_frame(sock), (130, b'after'))

    def test_slow_reader_and_handshake_do_not_block_worker(self):
        _, port = self.start(['--ws-buffer', '4096', '--socket-buffer', '4096', '--ws-write-ms', '300', '--ws-connect-ms', '300'])
        slow = self.connect(port, '/flood', receive_buffer=4096)
        stalled = socket.create_connection(('127.0.0.1', port), 2)
        self.sockets.append(stalled)
        stalled.settimeout(2)
        stalled.sendall(b'GET /game/stall HTTP/1.1\r\nHost: localhost\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n\r\n')
        normal = self.connect(port)
        for _ in range(10):
            normal.sendall(frame(b'play'))
            self.assertEqual(recv_frame(normal), (130, b'play'))
            self.asset(port)
            time.sleep(.05)
        self.assertTrue(stalled.recv(4096).startswith(b'HTTP/1.1 502'))
        time.sleep(.6)
        # A blocked reader is closed on the write deadline; buffered kernel bytes
        # may still be delivered before EOF, but memory remains bounded.
        slow.settimeout(2)
        total = 0
        while chunk := slow.recv(65536):
            total += len(chunk)
        self.assertLess(total, 8 * 1024 * 1024)

    def test_idle_timeout_and_graceful_drain(self):
        p, port = self.start(['--ws-idle-ms', '250', '--drain-ms', '700'])
        sock = self.connect(port)
        time.sleep(.45)
        self.assertEqual(sock.recv(1), b'')
        sock = self.connect(port)
        p.send_signal(signal.SIGTERM)
        time.sleep(.1)
        sock.sendall(frame(b'last update'))
        self.assertEqual(recv_frame(sock), (130, b'last update'))
        p.wait(timeout=2)
        self.assertEqual(p.returncode, 0)

    def test_default_idle_survives_old_ten_second_cutoff(self):
        _, port = self.start()
        sock = self.connect(port)
        time.sleep(10.3)
        sock.sendall(frame(b'still here'))
        self.assertEqual(recv_frame(sock), (130, b'still here'))

    def test_tls_slow_reader_does_not_block_other_sessions(self):
        if not self.tls:
            self.skipTest('TLS unavailable in Intel macOS build')
        _, port = self.start(['--ws-buffer', '4096', '--socket-buffer', '4096', '--ws-write-ms', '300'], tls=True)
        slow = self.connect(port, '/flood', tls=True, receive_buffer=4096)
        normal = self.connect(port, tls=True)
        for _ in range(20):
            normal.sendall(frame(b'play'))
            self.assertEqual(recv_frame(normal), (130, b'play'))
            time.sleep(.05)
        # macOS can autotune the TCP receive window during the TLS handshake.
        # Allow that window to fill before expecting the no-progress deadline.
        time.sleep(3)
        total = 0
        while chunk := slow.recv(65536): total += len(chunk)
        self.assertLess(total, 8 * 1024 * 1024)

    def test_reconnect_burst(self):
        _, port = self.start(['--connections', '128'])
        for _ in range(5):
            with concurrent.futures.ThreadPoolExecutor(max_workers=16) as pool:
                connections = list(pool.map(lambda _: self.connect(port), range(32)))
            for sock in connections:
                sock.sendall(frame(b'bye', opcode=8)); sock.close()
            time.sleep(.03)
        self.asset(port)


if __name__ == '__main__':
    unittest.main()