#!/usr/bin/env python3
"""The Debug Console of a Debug Build, over Wi-Fi (TCP 2323, see ADR 0004).

Usage: scripts/rdbg.py [-H host] [-b] [command ...]
  no command   interactive: type commands, see every console line live; Ctrl-D or `quit` leaves
  command      runs it and prints what follows, until the console has been quiet for a moment
  -H host      the device's IP (Settings > Firmware), default $RORO_OTA_HOST
  -b           also print the backlog the device sends on connecting (boot messages and so on)
The token is read from ~/.config/roro9stack/debug-token (made by the first build).
"""
import os
import select
import socket
import sys
import time

PORT = 2323
TOKEN_FILE = os.path.expanduser("~/.config/roro9stack/debug-token")


def read_until(sock, marker, timeout):
    """Everything up to and including `marker`, or None on a timeout or a closed connection."""
    sock.settimeout(timeout)
    data = b""
    while marker not in data:
        try:
            chunk = sock.recv(4096)
        except socket.timeout:
            return None
        if not chunk:
            return None
        data += chunk
    return data


def read_until_quiet(sock, quiet, out):
    """Copies everything the device sends to `out` until nothing has arrived for `quiet` seconds."""
    sock.settimeout(quiet)
    while True:
        try:
            data = sock.recv(4096)
        except socket.timeout:
            return True
        if not data:
            return False
        if out:
            out.write(data.decode(errors="replace"))
            out.flush()


def interactive(sock):
    sock.setblocking(False)
    while True:
        ready, _, _ = select.select([sock, sys.stdin], [], [])
        if sock in ready:
            data = sock.recv(4096)
            if not data:
                print("\n(connection closed)")
                return
            sys.stdout.write(data.decode(errors="replace"))
            sys.stdout.flush()
        if sys.stdin in ready:
            line = sys.stdin.readline()
            if not line:
                return
            sock.sendall(line.encode())


def main():
    args = sys.argv[1:]
    host, backlog = os.environ.get("RORO_OTA_HOST"), False
    while args and args[0].startswith("-"):
        flag = args.pop(0)
        if flag == "-H" and args:
            host = args.pop(0)
        elif flag == "-b":
            backlog = True
        else:
            sys.exit(__doc__)
    if not host:
        sys.exit("No device: pass -H <ip> or set RORO_OTA_HOST\n\n" + __doc__)
    token = open(TOKEN_FILE).read().strip()

    with socket.create_connection((host, PORT), timeout=10) as sock:
        sock.sendall((token + "\n").encode())
        # The device may still be finishing a previous client: wait for this connection's banner.
        banner = read_until(sock, b"Backlog follows.\n", 15)
        if banner is None:
            sys.exit("device: no banner (wrong token, or another client is connected)")
        show = sys.stdout if backlog or not args else None
        if show:
            show.write(banner.decode(errors="replace"))
        read_until_quiet(sock, 0.5, show)  # the backlog
        if not args:
            return interactive(sock)
        command = " ".join(args)
        sock.sendall((command + "\n").encode())
        # The main loop echoes "> command" when it runs it; the reply follows.
        echoed = read_until(sock, f"> {command}\n".encode(), 15)
        if echoed is None:
            sys.exit("device: the command never ran")
        sys.stdout.write(echoed.decode(errors="replace").split(f"> {command}\n", 1)[1])
        try:
            read_until_quiet(sock, 1.5, sys.stdout)
        except BrokenPipeError:  # e.g. piped into head
            pass


if __name__ == "__main__":
    main()