#!/usr/bin/env python3
"""WebSocket PTY client — like SSH but over WebSocket.

Puts your terminal in raw mode so TUIs (vim, nano, htop) work.
"""

import asyncio
import json
import os
import signal
import sys
import termios
import threading
import tty

import websockets


async def main():
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} wss://host[:port]", file=sys.stderr)
        sys.exit(1)

    url = sys.argv[1]
    if not url.startswith(("ws://", "wss://")):
        url = f"wss://{url}"

    fd = sys.stdin.fileno()
    old = termios.tcgetattr(fd)
    tty.setraw(fd)

    cols, rows = os.get_terminal_size()
    stdin_queue: asyncio.Queue = asyncio.Queue()

    def stdin_reader():
        try:
            while True:
                data = os.read(fd, 65536)
                if not data:
                    break
                stdin_queue.put_nowait(data)
        except Exception:
            pass

    threading.Thread(target=stdin_reader, daemon=True).start()

    def on_resize(signum, frame):
        nonlocal cols, rows
        try:
            nc, nr = os.get_terminal_size()
            if nc != cols or nr != rows:
                cols, rows = nc, nr
        except Exception:
            pass

    signal.signal(signal.SIGWINCH, on_resize)

    try:
        async with websockets.connect(url, max_size=2**20, ping_interval=None) as ws:
            # Send initial resize
            await ws.send(json.dumps({"rows": rows, "cols": cols}))

            async def sender():
                last_c, last_r = cols, rows
                while True:
                    if cols != last_c or rows != last_r:
                        await ws.send(json.dumps({"rows": rows, "cols": cols}))
                        last_c, last_r = cols, rows
                    try:
                        data = await asyncio.wait_for(stdin_queue.get(), timeout=0.1)
                        await ws.send(data)  # bytes → binary frame
                    except asyncio.TimeoutError:
                        continue

            async def receiver():
                async for data in ws:
                    if isinstance(data, str):
                        continue  # skip text frames (resize acks etc.)
                    sys.stdout.buffer.write(data)
                    sys.stdout.buffer.flush()

            await asyncio.gather(sender(), receiver())
    finally:
        termios.tcsetattr(fd, termios.TCSADRAIN, old)
        print()  # newline after raw mode


if __name__ == "__main__":
    asyncio.run(main())