#!/usr/bin/env python3
"""
deploy/smoke/ws_handshake.py — minimal WS client for smoke scripts (T8 factoring).

Hand-rolled RFC 6455 frame encode/decode + the Twilio Media Streams connected/start
handshake sequence, decoupled from `docker cp` / container orchestration so both
the all-in-one smoke (`allinone_smoke.py`) + the compose smoke (`compose_smoke.sh`,
via `python3 ws_handshake.py --host ... --port ... --ca-pem-file ...`) can drive
one trunk WS handshake + assert the EventSink wrote a lifecycle event.

No third-party deps — Python 3 + ssl + socket + json (same posture as allinone_smoke).
"""

import argparse
import json
import os
import socket
import ssl
import struct
import sys


def make_ws_frame(payload: bytes, opcode: int = 0x1) -> bytes:
    """Encode a client->server WS frame. opcode 0x1=text, 0x2=binary.
    Client frames MUST be masked (RFC 6455 §5.3)."""
    mask = os.urandom(4)
    masked = bytes(b ^ mask[i % 4] for i, b in enumerate(payload))
    header = bytearray([0x80 | opcode])  # FIN + opcode
    n = len(payload)
    if n < 126:
        header.append(0x80 | n)
    elif n < 65536:
        header.append(0x80 | 126)
        header.extend(struct.pack(">H", n))
    else:
        header.append(0x80 | 127)
        header.extend(struct.pack(">Q", n))
    header.extend(mask)
    return bytes(header) + masked


def read_ws_frame(sock) -> tuple:
    """Read one WS frame from the server. Returns (opcode, payload).
    Server frames are NOT masked (RFC 6455 §5.1).
    Returns (-1, b"") on EOF."""
    h1 = sock.recv(1)
    if not h1:
        return -1, b""
    fin_opcode = h1[0]
    h2 = sock.recv(1)[0]
    masked = bool(h2 & 0x80)
    length = h2 & 0x7F
    if length == 126:
        length = struct.unpack(">H", sock.recv(2))[0]
    elif length == 127:
        length = struct.unpack(">Q", sock.recv(8))[0]
    if masked:
        mask = sock.recv(4)
    payload = b""
    while len(payload) < length:
        chunk = sock.recv(length - len(payload))
        if not chunk:
            break
        payload += chunk
    if masked:
        payload = bytes(b ^ mask[i % 4] for i, b in enumerate(payload))
    return fin_opcode & 0x0F, payload


def open_wss(host: str, port: int, ca_pem: bytes, path: str = "/twilio/media-stream") -> ssl.SSLSocket:
    """Open a wss:// WS handshake to the given path. `ca_pem` is the raw PEM bytes
    (PEM, NOT a tar archive — the caller is responsible for extracting it from
    wherever it came from). `check_hostname=False` because Caddy's local_certs
    issues a cert for the SNI the operator set, not for `localhost`."""
    sock = socket.create_connection((host, port), timeout=5)
    ctx = ssl.create_default_context(cafile=None)
    ctx.load_verify_locations(cadata=ca_pem.decode("ascii"))
    ctx.check_hostname = False
    ssl_sock = ctx.wrap_socket(sock, server_hostname=host)

    key = os.urandom(16).hex()
    handshake = (
        f"GET {path} HTTP/1.1\r\n"
        f"Host: {host}:{port}\r\n"
        f"Upgrade: websocket\r\n"
        f"Connection: Upgrade\r\n"
        f"Sec-WebSocket-Key: {key}\r\n"
        f"Sec-WebSocket-Version: 13\r\n"
        f"\r\n"
    )
    ssl_sock.sendall(handshake.encode("ascii"))
    buf = b""
    while b"\r\n\r\n" not in buf:
        chunk = ssl_sock.recv(4096)
        if not chunk:
            raise RuntimeError("WS upgrade response closed prematurely")
        buf += chunk
    if b"101 Switching Protocols" not in buf:
        raise RuntimeError(f"WS upgrade did not return 101: {buf[:300]!r}")
    return ssl_sock


def send_twilio_handshake(sock) -> None:
    """Send the Twilio Media Streams connected + start frames (the same
    JSON sequence crates/rutster-trunk/tests/ws_ping.rs uses against the
    in-process router — here against the real containerized edge->FOB path)."""
    connected = json.dumps({
        "event": "connected",
        "protocol": "twilio-media-stream",
        "version": "1.0.0",
    })
    start = json.dumps({
        "event": "start",
        "start": {"streamSid": "MZsmoke", "callSid": "CAsmoke"},
    })
    sock.sendall(make_ws_frame(connected.encode("utf-8")))
    sock.sendall(make_ws_frame(start.encode("utf-8")))


def main() -> int:
    """CLI entry — used by compose_smoke.sh. Drives one connected+start
    handshake, drains a handful of server frames, exits 0 on success.

    Designed to be invoked from a bash smoke script that already extracted
    the root cert PEM to a file (the compose smoke does this via
    `docker compose exec -T caddy cat ...`)."""
    p = argparse.ArgumentParser(description="Minimal WS handshake client for smoke tests.")
    p.add_argument("--host", default="localhost")
    p.add_argument("--port", type=int, default=443)
    p.add_argument("--ca-pem-file", required=True,
                   help="Path to the Caddy internal CA root cert PEM (NOT a tar archive).")
    p.add_argument("--path", default="/twilio/media-stream")
    args = p.parse_args()

    with open(args.ca_pem_file, "rb") as f:
        ca_pem = f.read()

    sock = open_wss(args.host, args.port, ca_pem, args.path)
    send_twilio_handshake(sock)
    # Drain a handful of frames (mark/outbound/etc.) — we just need the
    # EventSink to fire on_event once; the XLEN assertion in the bash script
    # is the real gate, not the frame count here.
    import time
    deadline = time.time() + 2.0
    received = 0
    while time.time() < deadline:
        sock.settimeout(0.5)
        try:
            op, _ = read_ws_frame(sock)
        except (socket.timeout, OSError):
            break
        if op in (0x1, 0x2):
            received += 1
        elif op == -1:
            break
    sock.close()
    print(f"[ws_handshake] sent connected+start; received {received} frames back")
    return 0


if __name__ == "__main__":
    sys.exit(main())
