#!/usr/bin/env python3
"""Capture a screenshot from the watch over USB serial (Windows).

The Unix `screenshot.sh` shells out to ffmpeg; this is a self-contained Windows
equivalent that needs only pyserial + Pillow (both in the daemon venv). The
firmware's `screenshot` serial command dumps the LVGL framebuffer as raw RGB565LE
between `SCREENSHOT_START w h size` and `SCREENSHOT_END`.

Usage:
    python tools/screenshot_win.py [COM_PORT] [out.png]      (defaults: COM7 shot.png)
"""
import array
import sys
import time

import serial
from PIL import Image

port_name = sys.argv[1] if len(sys.argv) > 1 else "COM7"
out = sys.argv[2] if len(sys.argv) > 2 else "shot.png"

port = serial.Serial(port_name, 115200, timeout=3)
time.sleep(0.3)
port.reset_input_buffer()
port.write(b"screenshot\n")
port.flush()

w = h = raw_size = 0
deadline = time.time() + 15
while time.time() < deadline:
    line = port.readline().decode("utf-8", errors="replace").strip()
    if line.startswith("SCREENSHOT_START"):
        _, sw, sh, ss = line.split()
        w, h, raw_size = int(sw), int(sh), int(ss)
        break
    if line in ("SCREENSHOT_ERR", "SCREENSHOT_UNSUPPORTED"):
        sys.exit(f"device error: {line}")
if not raw_size:
    sys.exit("no SCREENSHOT_START (is the device booted and on COM port?)")

data = b""
while len(data) < raw_size:
    chunk = port.read(min(8192, raw_size - len(data)))
    if not chunk:
        sys.exit(f"timeout: got {len(data)} of {raw_size} bytes")
    data += chunk
port.close()

# RGB565 little-endian -> RGB888. array 'H' is native (LE on x86) == RGB565LE.
px = array.array("H")
px.frombytes(data)
rgb = bytearray(len(px) * 3)
for i, v in enumerate(px):
    r = (v >> 11) & 0x1F
    g = (v >> 5) & 0x3F
    b = v & 0x1F
    rgb[i * 3] = (r << 3) | (r >> 2)
    rgb[i * 3 + 1] = (g << 2) | (g >> 4)
    rgb[i * 3 + 2] = (b << 3) | (b >> 2)
Image.frombytes("RGB", (w, h), bytes(rgb)).save(out)
print(f"saved {out} {w}x{h} ({len(data)} bytes)")