#!/usr/bin/env python3
"""Network screenshot of a running VT1200, using rdebug `peek` over TCP 771.

No debugger, no ROM patch, no agent on the terminal: it reads the framebuffer
out of video RAM one `peek` at a time and writes a PNG.

The memory layout is the one the harness renderer uses
(`render_active_line()` in harness/mame/vt1200.cpp):

  * base   bit 0x04000000   (video RAM)
  * pitch  1024 bits = 128 bytes per row
  * bit order  LSB-first within each byte/word
  * active  1024 x 1024 pixels; on the emulated screen it appears at
            x 192..1215, y 36..900 (MAME-derived blanking; see --screen)

Examples:

  # the raw 1024x1024 OS canvas
  python3 dwt/screenshot.py --host 192.168.99.2 --out shot.png

  # compose the emulator's visible 1280x901 screen, for comparison with MAME dumps
  python3 dwt/screenshot.py --host 192.168.99.2 --screen --out screen.png

  # only the top of the canvas, and keep the raw framebuffer too
  python3 dwt/screenshot.py --host 192.168.99.2 --rows 0:256 \\
      --out top.png --raw vram.bin

  # a second grab 1 s later (tearing check)
  python3 dwt/screenshot.py --host 192.168.99.2 --repeat 2 --interval 1.0

The peek limit is 1024 bytes, i.e. exactly 8 rows per request, so a full
1024-row grab is 128 round trips.
"""
from __future__ import annotations

import argparse
import struct
import sys
import time
import zlib

import exploit

VRAM_BASE_BITS = 0x04000000
PITCH_BITS = 1024
ROW_BYTES = PITCH_BITS // 8          # 128
WIDTH = 1024
HEIGHT = 1024
CHUNK = exploit.MAX_XFER             # 1024 bytes = 8 rows
ROWS_PER_CHUNK = CHUNK // ROW_BYTES  # 8

# Defaults for the emulated screen; the real unit's blanking may differ, which
# is why the raw canvas is the primary artifact and --screen is opt-in.
SCREEN_W = 1280
SCREEN_H = 901
OFFSET_X = 192
OFFSET_Y = 36


def parse_rows(text: str) -> tuple[int, int]:
    if ":" in text:
        a, b = text.split(":", 1)
        start, end = int(a, 0), int(b, 0)
    else:
        start, end = 0, int(text, 0)
    if not (0 <= start < end <= HEIGHT):
        raise argparse.ArgumentTypeError(
            f"rows must satisfy 0 <= start < end <= {HEIGHT}")
    return start, end


def decode(data: bytes, width: int = WIDTH) -> bytearray:
    """LSB-first bits -> 8-bit grayscale (bit 1 = white)."""
    rows = len(data) // ROW_BYTES
    out = bytearray(rows * width)
    for r in range(rows):
        src = r * ROW_BYTES
        dst = r * width
        for x in range(width):
            v = (data[src + (x >> 3)] >> (x & 7)) & 1
            out[dst + x] = 255 if v else 0
    return out


def compose_screen(gray: bytearray, rows: int, sw: int, sh: int,
                   ox: int, oy: int) -> bytearray:
    canvas = bytearray(sw * sh)
    for r in range(rows):
        y = r + oy
        if 0 <= y < sh:
            xs = max(0, ox)
            xe = min(sw, ox + WIDTH)
            if xe > xs:
                canvas[y * sw + xs:y * sw + xe] = \
                    gray[r * WIDTH + (xs - ox):r * WIDTH + (xe - ox)]
    return canvas


def write_png(path: str, width: int, height: int, gray: bytes) -> None:
    raw = bytearray()
    for y in range(height):
        raw.append(0)                      # filter: none
        raw += gray[y * width:(y + 1) * width]
    comp = zlib.compress(bytes(raw), 9)

    def chunk(tag: bytes, data: bytes) -> bytes:
        return (struct.pack(">I", len(data)) + tag + data +
                struct.pack(">I", zlib.crc32(tag + data) & 0xFFFFFFFF))

    png = b"\x89PNG\r\n\x1a\n"
    png += chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 0, 0, 0, 0))
    png += chunk(b"IDAT", comp)
    png += chunk(b"IEND", b"")
    with open(path, "wb") as f:
        f.write(png)


def grab(rd: exploit.RDebug, row_start: int, row_end: int,
         quiet: bool) -> bytes:
    start = row_start * ROW_BYTES
    end = row_end * ROW_BYTES
    data = bytearray()
    pos = start
    nchunks = (end - start + CHUNK - 1) // CHUNK
    done = 0
    while pos < end:
        n = min(CHUNK, end - pos)
        data += rd.peek(VRAM_BASE_BITS + pos * 8, n)
        pos += n
        done += 1
        if not quiet and nchunks > 8 and (done % 16 == 0 or done == nchunks):
            print(f"\r[.] {done}/{nchunks} chunks", end="", file=sys.stderr,
                  flush=True)
    if not quiet and nchunks > 8:
        print(file=sys.stderr)
    return bytes(data)


def one(rd: exploit.RDebug, a: argparse.Namespace, tag: str) -> None:
    data = grab(rd, a.rows[0], a.rows[1], a.quiet)
    rows = len(data) // ROW_BYTES
    gray = decode(data)
    if a.invert:
        gray = bytearray(255 - v for v in gray)

    out = a.out
    if a.repeat > 1:
        root, dot, ext = out.rpartition(".")
        out = f"{root}-{tag}.{ext}" if dot else f"{out}-{tag}"
    if a.screen:
        canvas = compose_screen(gray, rows, a.screen_width, a.screen_height,
                                a.offset_x, a.offset_y)
        write_png(out, a.screen_width, a.screen_height, canvas)
        print(f"[+] {out}: {a.screen_width}x{a.screen_height} "
              f"(canvas {WIDTH}x{HEIGHT}, rows {a.rows[0]}..{a.rows[1]-1})")
    else:
        write_png(out, WIDTH, rows, gray)
        print(f"[+] {out}: {WIDTH}x{rows} (rows {a.rows[0]}..{a.rows[1]-1})")

    if a.raw:
        rp = a.raw
        if a.repeat > 1:
            root, dot, ext = rp.rpartition(".")
            rp = f"{root}-{tag}.{ext}" if dot else f"{rp}-{tag}"
        with open(rp, "wb") as f:
            f.write(data)
        print(f"[+] {rp}: {len(data)} bytes of VRAM")


def main(argv: list[str]) -> int:
    ap = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--host", default="192.168.99.2")
    ap.add_argument("--port", type=int, default=771)
    ap.add_argument("--timeout", type=float, default=8.0)
    ap.add_argument("--out", default="shot.png")
    ap.add_argument("--rows", type=parse_rows, default=(0, HEIGHT),
                    metavar="A:B", help=f"row range to read (default 0:{HEIGHT})")
    ap.add_argument("--screen", action="store_true",
                    help="compose the visible screen instead of the raw canvas")
    ap.add_argument("--screen-width", type=int, default=SCREEN_W)
    ap.add_argument("--screen-height", type=int, default=SCREEN_H)
    ap.add_argument("--offset-x", type=int, default=OFFSET_X)
    ap.add_argument("--offset-y", type=int, default=OFFSET_Y)
    ap.add_argument("--invert", action="store_true", help="white <-> black")
    ap.add_argument("--raw", metavar="FILE", help="also write the raw VRAM bytes")
    ap.add_argument("--repeat", type=int, default=1, help="grab N times")
    ap.add_argument("--interval", type=float, default=0.0,
                    help="seconds between repeats")
    ap.add_argument("--quiet", action="store_true")
    ap.add_argument("--verbose", action="store_true")
    a = ap.parse_args(argv[1:])

    rd = exploit.RDebug(a.host, a.port, a.timeout, a.verbose)
    rd.connect()
    try:
        rd.liveness(7)                 # cheap reachability check
        for i in range(max(1, a.repeat)):
            if i:
                time.sleep(a.interval)
            one(rd, a, time.strftime("%H%M%S") if a.repeat > 1 else "1")
    finally:
        rd.close()
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main(sys.argv))
    except exploit.ProtocolError as exc:
        print(f"error: {exc}", file=sys.stderr)
        sys.exit(1)
    except OSError as exc:
        print(f"error: {exc}", file=sys.stderr)
        sys.exit(2)