#!/usr/bin/env python3
"""Capture a screenshot from the DeskLock device over USB.

The firmware (SCREENSHOT_ENABLE) watches the USB-serial-JTAG RX for a trigger
byte and streams the current screen as a text header
  ###SHOT_BEGIN w=.. h=.. fmt=rgb565le bytes=N crc=0xXXXX bin=1###\n
followed by exactly N raw RGB565-LE bytes, then \n###SHOT_END###. We send the
trigger, read the payload, verify the CRC, and write a PNG.

Usage (run in the dialout group, e.g. under `sg dialout -c '...'`):
    python3 device_shot.py [out.png] [--overlay] [--port /dev/ttyACM0]
  --overlay  send 'o' (force the tap controls overlay) instead of 's' (current screen)
"""
import os, sys, termios, select, time, zlib, re

port = "/dev/ttyACM0"
out = "shot.png"
trig = b"s"
args = sys.argv[1:]
while args:
    a = args.pop(0)
    if a == "--overlay": trig = b"o"
    elif a == "--port": port = args.pop(0)
    elif not a.startswith("-"): out = a
    else: sys.stderr.write(f"[warn] ignoring {a}\n")

ATTEMPTS = 4
PER_TRY = 8.0
begin_re = re.compile(
    rb"###SHOT_BEGIN w=(\d+) h=(\d+) fmt=(\S+) bytes=(\d+) crc=0x([0-9a-fA-F]+) bin=1###\n")

fd = os.open(port, os.O_RDWR | os.O_NOCTTY | os.O_NONBLOCK)
a = termios.tcgetattr(fd)
a[0] = 0; a[1] = 0; a[3] = 0            # raw: no CR/NL translation, no canon/echo
a[4] = termios.B115200; a[5] = termios.B115200
termios.tcsetattr(fd, termios.TCSANOW, a)


def drain(sec=0.3):
    t = time.time()
    while time.time() - t < sec:
        r, _, _ = select.select([fd], [], [], 0.05)
        if r:
            try: os.read(fd, 65536)
            except BlockingIOError: pass


def capture_once():
    drain(0.3)                          # clear any in-flight tail from a prior dump
    buf = bytearray(); hdr = None; t0 = time.time(); last_trig = 0.0
    while time.time() - t0 < PER_TRY:
        now = time.time()
        if hdr is None and now - last_trig >= 2.0:
            try: os.write(fd, trig)
            except OSError: pass
            last_trig = now
        r, _, _ = select.select([fd], [], [], 0.2)
        if r:
            try: chunk = os.read(fd, 65536)
            except BlockingIOError: chunk = b""
            if chunk: buf += chunk
        if hdr is None:
            m = begin_re.search(buf)
            if m:
                hdr = (int(m.group(1)), int(m.group(2)), int(m.group(4)), int(m.group(5), 16))
                buf = bytearray(buf[m.end():])      # everything after header = raw payload
        if hdr is not None and len(buf) >= hdr[2]:
            return (hdr[0], hdr[1], hdr[2], hdr[3], bytes(buf[:hdr[2]]))
    return None


result = None
for attempt in range(1, ATTEMPTS + 1):
    res = capture_once()
    if res is None:
        sys.stderr.write(f"[try {attempt}] no frame\n"); continue
    w, h, nbytes, crc_dev, raw = res
    crc_host = zlib.crc32(raw) & 0xffffffff
    ok = crc_host == crc_dev
    sys.stderr.write(f"[try {attempt}] {w}x{h} {len(raw)}B "
                     f"crc dev=0x{crc_dev:08x} host=0x{crc_host:08x} {'OK' if ok else 'RETRY'}\n")
    if ok:
        result = (w, h, raw); break
os.close(fd)

if result is None:
    sys.stderr.write("[fail] no clean frame — SCREENSHOT_ENABLE flashed? device up?\n")
    sys.exit(2)

w, h, raw = result
from PIL import Image
img = Image.new("RGB", (w, h)); px = img.load(); i = 0
for y in range(h):
    for x in range(w):
        v = raw[i] | (raw[i+1] << 8); i += 2
        r5 = (v >> 11) & 0x1f; g6 = (v >> 5) & 0x3f; b5 = v & 0x1f
        px[x, y] = ((r5 << 3) | (r5 >> 2), (g6 << 2) | (g6 >> 4), (b5 << 3) | (b5 >> 2))
img.save(out)
sys.stderr.write(f"[ok] wrote {out} ({w}x{h})\n")
print(out)