#!/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
  get <card path> [file]   copy a file from the SD card
  put <file> [card path]   copy a file to the SD card (default /updates/<name>), checked with SHA-256
  screenshot [file.png]    what the screen shows (default screen-<date>.png), at 2x
  reset                    restart at once, even if the main loop is stuck
The token is read from ~/.config/roro9stack/debug-token (made by the first build).
"""
import hashlib
import os
import re
import select
import struct
import subprocess
import zlib
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 binary(sock, command, tag):
    """Sends a binary command; returns (header words, payload) or exits with the device's error."""
    sock.sendall((command + "\n").encode())
    buf = b""
    sock.settimeout(30)
    pattern = re.compile(tag.encode() + rb": (data|ready|rgb332|error)([^\n]*)\n")
    while not (m := pattern.search(buf)):
        chunk = sock.recv(65536)
        if not chunk:
            sys.exit("device: closed the connection")
        buf += chunk
    if m.group(1) == b"error":
        sys.exit(f"device: {tag}: error{m.group(2).decode()}")
    return m.group(1).decode(), m.group(2).decode().split(), buf[m.end():]


def receive(sock, have, size):
    while len(have) < size:
        chunk = sock.recv(65536)
        if not chunk:
            sys.exit(f"device: the connection closed after {len(have)} of {size} bytes")
        have += chunk
    return have[:size]


def get(sock, remote, local):
    _, words, rest = binary(sock, f"get {remote}", "get")
    start, size = time.time(), int(words[0])
    data = receive(sock, rest, size)
    local = local or os.path.basename(remote)
    open(local, "wb").write(data)
    print(f"saved {local} ({size} bytes, {size / 1024 / max(time.time() - start, 0.001):.0f} KB/s)")


def put(sock, local, remote):
    data = open(local, "rb").read()
    remote = remote or "/updates/" + os.path.basename(local)
    digest = hashlib.sha256(data).hexdigest()
    binary(sock, f"put {remote} {len(data)} {digest}", "put")
    start = time.time()
    sock.sendall(data)
    answer = read_until(sock, b"\n", 30) or b"put: error no answer"
    answer = answer.decode(errors="replace").strip().splitlines()[-1]
    print(f"device: {answer} ({len(data) / 1024 / max(time.time() - start, 0.001):.0f} KB/s)")
    if " error " in answer:
        sys.exit(1)


def png(path, width, height, rgb, scale):
    rows = b""
    for y in range(height):
        row = b"".join(rgb[(y * width + x) * 3:(y * width + x) * 3 + 3] * scale for x in range(width))
        rows += (b"\x00" + row) * scale
    def chunk(kind, body):
        return struct.pack(">I", len(body)) + kind + body + struct.pack(">I", zlib.crc32(kind + body))
    header = struct.pack(">IIBBBBB", width * scale, height * scale, 8, 2, 0, 0, 0)
    with open(path, "wb") as f:
        f.write(b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", header) + chunk(b"IDAT", zlib.compress(rows)) + chunk(b"IEND", b""))


def screenshot(sock, path):
    _, words, rest = binary(sock, "screenshot", "screenshot")
    width, height = int(words[0]), int(words[1])
    pixels = receive(sock, rest, width * height)
    # RGB332, as M5GFX stores an 8-bit sprite: RRRGGGBB.
    rgb = b"".join(bytes(((v >> 5) * 255 // 7, ((v >> 2) & 7) * 255 // 7, (v & 3) * 255 // 3)) for v in pixels)
    path = path or time.strftime("screen-%Y%m%d-%H%M%S.png")
    png(path, width, height, rgb, 2)
    print(f"saved {path} ({width * 2}x{height * 2})")


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[0] == "get" and len(args) >= 2:
            return get(sock, args[1], args[2] if len(args) > 2 else None)
        if args[0] == "put" and len(args) >= 2:
            return put(sock, args[1], args[2] if len(args) > 2 else None)
        if args[0] == "screenshot":
            return screenshot(sock, args[1] if len(args) > 1 else None)
        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?)")