#!/usr/bin/env python3
"""Generic stdlib-only WebSocket bridge: runs a command, streams its stdout to
every connected WebSocket client as binary frames. One-directional (server ->
browser) by design, so it never parses masked client frames beyond detecting a
closed connection - built for low-latency PCM audio streaming, but the command
run is arbitrary.

Usage: pcm_ws_bridge.py <port> <command> [args...]
"""

import base64
import hashlib
import socket
import subprocess
import sys
import threading

WS_MAGIC = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

clients = []
clients_lock = threading.Lock()


def handshake(conn):
    request = b""
    while b"\r\n\r\n" not in request:
        chunk = conn.recv(4096)
        if not chunk:
            return False
        request += chunk
    key = None
    for line in request.split(b"\r\n")[1:]:
        if line.lower().startswith(b"sec-websocket-key:"):
            key = line.split(b":", 1)[1].strip()
            break
    if not key:
        return False
    accept = base64.b64encode(hashlib.sha1(key + WS_MAGIC.encode()).digest()).decode()
    conn.sendall(
        (
            "HTTP/1.1 101 Switching Protocols\r\n"
            "Upgrade: websocket\r\n"
            "Connection: Upgrade\r\n"
            f"Sec-WebSocket-Accept: {accept}\r\n\r\n"
        ).encode()
    )
    return True


def frame(data):
    length = len(data)
    if length <= 125:
        header = bytes([0x82, length])
    elif length <= 0xFFFF:
        header = bytes([0x82, 126]) + length.to_bytes(2, "big")
    else:
        header = bytes([0x82, 127]) + length.to_bytes(8, "big")
    return header + data


def handle_client(conn):
    if not handshake(conn):
        conn.close()
        return
    with clients_lock:
        clients.append(conn)
    try:
        while conn.recv(4096):
            pass  # not expecting client messages, just watching for disconnect
    except OSError:
        pass
    finally:
        with clients_lock:
            if conn in clients:
                clients.remove(conn)
        conn.close()


def broadcast_loop(proc):
    while True:
        chunk = proc.stdout.read(4096)
        if not chunk:
            break
        packet = frame(chunk)
        with clients_lock:
            dead = []
            for c in clients:
                try:
                    c.sendall(packet)
                except OSError:
                    dead.append(c)
            for c in dead:
                clients.remove(c)


def main():
    port = int(sys.argv[1])
    cmd = sys.argv[2:]
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE)
    threading.Thread(target=broadcast_loop, args=(proc,), daemon=True).start()

    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server.bind(("0.0.0.0", port))
    server.listen(5)
    while True:
        conn, _ = server.accept()
        threading.Thread(target=handle_client, args=(conn,), daemon=True).start()


if __name__ == "__main__":
    main()