#!/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)

Commands handled here as well as on the device:
  crash             the last crash, with its backtrace decoded (scripts/decode_backtrace.sh)
  coredump [file]   fetch the core dump (default core-<date>.bin) and decode it with esp-coredump
The token is read from ~/.config/roro9stack/debug-token (made by the first build).
"""
import os
import re
import select
import subprocess
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()


SCRIPTS = os.path.dirname(os.path.abspath(__file__))


def run(sock, command, out=sys.stdout):
    """Sends one command and returns its reply (also copied to `out`)."""
    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")
    reply = echoed.decode(errors="replace").split(f"> {command}\n", 1)[1]

    class Tee:
        def write(self, text):
            nonlocal reply
            reply += text
            if out:
                out.write(text)

        def flush(self):
            if out:
                out.flush()

    if out:
        out.write(reply)
    try:
        read_until_quiet(sock, 1.5, Tee())
    except BrokenPipeError:  # e.g. piped into head
        pass
    return reply


def crash_firmware(reply):
    """The archived-ELF key for the crashed firmware: its ELF digest, or else its version."""
    sha = re.search(r"elf sha256 ([0-9a-f]{8,})", reply)
    if sha:
        return sha.group(1)
    version = re.search(r"last one in (\S+)", reply)
    return version.group(1) if version else None


def crash(sock):
    reply = run(sock, "crash")
    trace = re.search(r"backtrace((?: 0x[0-9a-f]+)+)", reply)
    key = crash_firmware(reply)
    if trace and key:
        print()
        subprocess.call([os.path.join(SCRIPTS, "decode_backtrace.sh"), key] + trace.group(1).split())


def coredump(sock, path):
    info = run(sock, "crash", out=None)
    sock.sendall(b"coredump get\n")
    # One buffer throughout: the header, the size and the first bytes often share a packet.
    buf = b""
    sock.settimeout(15)
    while b"coredump: none" not in buf and not re.search(rb"coredump: data (\d+)\n", buf):
        chunk = sock.recv(65536)
        if not chunk:
            sys.exit("device: closed the connection")
        buf += chunk
    header = re.search(rb"coredump: data (\d+)\n", buf)
    if not header:
        sys.exit("device: no core dump in flash")
    size = int(header.group(1))
    data = buf[header.end():]
    while len(data) < size:
        chunk = sock.recv(65536)
        if not chunk:
            sys.exit(f"device: the connection closed after {len(data)} of {size} bytes")
        data += chunk
    data = data[:size]
    path = path or time.strftime("core-%Y%m%d-%H%M%S.bin")
    open(path, "wb").write(data)
    print(f"saved {path} ({size} bytes)")
    key = crash_firmware(info)
    if key:
        subprocess.call([os.path.join(SCRIPTS, "decode_coredump.sh"), path, key])


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)
        if args == ["crash"]:
            return crash(sock)
        if args == ["reset"]:  # answered by the console's own task, not the main loop
            sock.sendall(b"reset\n")
            print((read_until(sock, b"restarting now\n", 10) or b"device: no answer").decode().strip().splitlines()[-1])
            return
        if args[0] == "coredump" and args[1:2] != ["erase"]:
            return coredump(sock, args[1] if len(args) > 1 else None)
        run(sock, " ".join(args))


if __name__ == "__main__":
    try:
        main()
    except (ConnectionResetError, BrokenPipeError):
        sys.exit("device: the connection dropped (restarting?)")