#!/usr/bin/env python3
"""Grab the radio's screen over USB and save it as a PNG.

Uses the CPS read command (area 6 = display buffer), which the firmware answers only
while the radio is NOT in modem mode - and, importantly, such requests do not switch it
into modem mode either. So the screen can be watched while the radio does its normal job.

Display is 160x128, RGB565, 2 bytes per pixel.
"""
import os, sys, termios, time, zlib, struct, argparse

WIDTH, HEIGHT = 160, 128
CPS_ACCESS_DISPLAY_BUFFER = 6
CHUNK = 1024


def open_port(path):
    fd = os.open(path, os.O_RDWR | os.O_NOCTTY)
    a = termios.tcgetattr(fd)
    a[0] = 0
    a[1] = 0
    a[2] = termios.CS8 | termios.CREAD | termios.CLOCAL
    a[3] = 0
    a[6][termios.VMIN] = 0
    a[6][termios.VTIME] = 20
    termios.tcsetattr(fd, termios.TCSANOW, a)
    termios.tcflush(fd, termios.TCIOFLUSH)
    return fd


def read_exact(fd, n, timeout=3.0):
    buf = b""
    deadline = time.time() + timeout
    while len(buf) < n and time.time() < deadline:
        try:
            part = os.read(fd, n - len(buf))
        except BlockingIOError:
            part = b""
        if part:
            buf += part
        else:
            time.sleep(0.003)
    return buf


def read_block(fd, area, addr, length):
    req = bytes([ord('R'), area,
                 (addr >> 24) & 0xFF, (addr >> 16) & 0xFF, (addr >> 8) & 0xFF, addr & 0xFF,
                 (length >> 8) & 0xFF, length & 0xFF])
    termios.tcflush(fd, termios.TCIFLUSH)
    os.write(fd, req)
    head = read_exact(fd, 3)
    if len(head) < 3 or head[0:1] != b'R':
        return None
    got = (head[1] << 8) | head[2]
    data = read_exact(fd, got)
    return data if len(data) == got else None


def write_png(path, rgb_rows):
    """Minimal PNG writer (stdlib only)."""
    raw = b"".join(b"\x00" + row for row in rgb_rows)

    def chunk(tag, data):
        c = struct.pack(">I", len(data)) + tag + data
        return c + 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, 2, 0, 0, 0))
    png += chunk(b"IDAT", zlib.compress(raw, 6))
    png += chunk(b"IEND", b"")
    open(path, "wb").write(png)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", default="/dev/ttyACM0")
    ap.add_argument("--out", default="/home/viktor/screen.png")
    args = ap.parse_args()

    fd = open_port(args.port)
    try:
        total = WIDTH * HEIGHT * 2
        data = b""
        while len(data) < total:
            n = min(CHUNK, total - len(data))
            part = read_block(fd, CPS_ACCESS_DISPLAY_BUFFER, len(data), n)
            if part is None:
                print("read failed at offset %d - is the radio in modem mode?" % len(data))
                return 2
            data += part
    finally:
        os.close(fd)

    rows = []
    for y in range(HEIGHT):
        row = bytearray()
        for x in range(WIDTH):
            i = (y * WIDTH + x) * 2
            # RGB565, little endian as stored by the display driver
            v = data[i] | (data[i + 1] << 8)
            r = ((v >> 11) & 0x1F) << 3
            g = ((v >> 5) & 0x3F) << 2
            b = (v & 0x1F) << 3
            row += bytes((r, g, b))
        rows.append(bytes(row))

    write_png(args.out, rows)
    nonblack = sum(1 for i in range(0, len(data), 2) if data[i] or data[i + 1])
    print("saved %s  (%d of %d pixels non-black)" % (args.out, nonblack, WIDTH * HEIGHT))
    return 0


if __name__ == "__main__":
    sys.exit(main())