#!/usr/bin/env python3
import argparse
import base64
import json
import re
import sys
import time

import websocket


def strip_ansi(data: bytes) -> str:
    text = data.decode("utf-8", "replace")
    # CSI / OSC / ESC cleanup, enough for tty captures.
    text = re.sub(r"\x1b\][^\x07]*(?:\x07|\x1b\\)", "", text)
    text = re.sub(r"\x1b\[[0-?]*[ -/]*[@-~]", "", text)
    text = re.sub(r"\x1b[()][A-Za-z0-9]", "", text)
    text = text.replace("\r", "\n")
    return text


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--url", default="ws://192.168.1.116/sudoku/ws")
    ap.add_argument("--user", default="admin")
    ap.add_argument("--password", default="babylove3")
    ap.add_argument("--cols", type=int, default=120)
    ap.add_argument("--rows", type=int, default=40)
    ap.add_argument("--seconds", type=float, default=8)
    ap.add_argument("--send", action="append", default=[], help="send terminal input string after connect")
    ap.add_argument("--send-file", help="read terminal input from this file and send after connect")
    ap.add_argument("--send-delay", type=float, default=0.0, help="seconds to wait before sending input")
    ap.add_argument("--quiet", action="store_true", help="suppress per-frame stderr logs")
    ap.add_argument("--stripped-only", action="store_true", help="print ANSI-stripped terminal text instead of raw bytes")
    ap.add_argument("--tail-lines", type=int, default=0, help="with --stripped-only, print only the last N lines")
    args = ap.parse_args()

    auth = base64.b64encode(f"{args.user}:{args.password}".encode()).decode()
    ws = websocket.create_connection(
        args.url,
        subprotocols=["tty"],
        header=[f"Authorization: Basic {auth}"],
        timeout=3,
    )
    if not args.quiet:
        print(f"[+] connected subprotocol={ws.subprotocol!r}", file=sys.stderr)

    init = json.dumps(
        {"AuthToken": "", "columns": args.cols, "rows": args.rows},
        separators=(",", ":"),
    ).encode()
    ws.send(init, opcode=websocket.ABNF.OPCODE_BINARY)

    send_items = list(args.send)
    if args.send_file:
        with open(args.send_file, "r", encoding="utf-8") as f:
            send_items.append(f.read())
    if send_items and args.send_delay:
        time.sleep(args.send_delay)
    for s in send_items:
        payload = b"0" + s.encode()
        ws.send(payload, opcode=websocket.ABNF.OPCODE_BINARY)

    raw_out = bytearray()
    end = time.time() + args.seconds
    while time.time() < end:
        try:
            frame = ws.recv_frame()
        except Exception as e:
            print(f"[!] recv: {e!r}", file=sys.stderr)
            break
        if frame.opcode == websocket.ABNF.OPCODE_PING:
            ws.pong(frame.data)
            if not args.quiet:
                print(f"[.] ping -> pong {len(frame.data)}", file=sys.stderr)
            continue
        if frame.opcode == websocket.ABNF.OPCODE_CLOSE:
            if not args.quiet:
                print(f"[!] close {frame.data!r}", file=sys.stderr)
            break
        data = frame.data
        if not data:
            continue
        cmd, body = data[:1], data[1:]
        if not args.quiet:
            print(f"[.] frame opcode={frame.opcode} cmd={cmd!r} len={len(body)}", file=sys.stderr)
        if cmd == b"0":
            raw_out += body
        elif cmd in (b"1", b"2"):
            if not args.quiet:
                print(body.decode("utf-8", "replace"), file=sys.stderr)
        else:
            if not args.quiet:
                print(f"[?] raw {data[:120]!r}", file=sys.stderr)
    try:
        ws.close()
    except Exception:
        pass
    if args.stripped_only:
        text = strip_ansi(bytes(raw_out))
        if args.tail_lines > 0:
            text = "\n".join(text.splitlines()[-args.tail_lines:]) + "\n"
        print(text, end="")
    else:
        sys.stdout.buffer.write(raw_out)
        sys.stdout.flush()
        print("\n--- stripped ---", file=sys.stderr)
        print(strip_ansi(bytes(raw_out)), file=sys.stderr)


if __name__ == "__main__":
    main()