#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.9"
# dependencies = [
#     "requests",
#     "numpy",
#     "pillow",
# ]
# ///
"""Fetch an LVGL device's screen via the LvScreenshot HTTP endpoint and save it as a PNG.

The firmware must expose the endpoint (register it with lvScreenshotAttachHttp, e.g. behind a
build flag). The device returns the raw RGB565 framebuffer plus X-Width / X-Height / X-Format
headers; this script converts it to RGB888 and writes a PNG.

Dependencies are declared inline (PEP 723) and installed automatically by uv -- no manual setup.

Usage (via uv):
    uv run screenshot.py [--host diystickynotes.local] [--out shot.png] [--swap-rb] [--byteswap]
    # the shebang also lets you run it directly: ./screenshot.py ...

--swap-rb   : swap red/blue if the panel uses BGR order (colours look inverted).
--byteswap  : byte-swap each 16-bit pixel if the panel stores RGB565 big-endian (garbled image).
"""
import argparse
import sys

import numpy as np
import requests
from PIL import Image


def main() -> int:
    """Fetch the screen, convert RGB565 -> RGB888, and write a PNG. Returns a process exit code."""
    ap = argparse.ArgumentParser(description="Capture an LVGL device screen as a PNG.")
    ap.add_argument("--host", default="diystickynotes.local", help="device hostname or IP")
    ap.add_argument("--out", default="shot.png", help="output PNG path")
    ap.add_argument("--path", default="/screenshot", help="endpoint path")
    ap.add_argument("--swap-rb", action="store_true", help="swap R/B (BGR panel)")
    ap.add_argument("--byteswap", action="store_true", help="byte-swap 16-bit pixels")
    ap.add_argument("--timeout", type=float, default=15.0, help="HTTP timeout seconds")
    args = ap.parse_args()

    url = f"http://{args.host}{args.path}"
    try:
        r = requests.get(url, timeout=args.timeout)
    except requests.RequestException as e:
        print(f"error: request to {url} failed: {e}", file=sys.stderr)
        return 1
    if r.status_code != 200:
        print(f"error: HTTP {r.status_code}: {r.text[:200]}", file=sys.stderr)
        return 1

    try:
        w = int(r.headers["X-Width"])
        h = int(r.headers["X-Height"])
        fmt = r.headers.get("X-Format", "RGB565")
    except (KeyError, ValueError) as e:
        print(f"error: missing/invalid dimension headers: {e}", file=sys.stderr)
        return 1
    if fmt != "RGB565":
        print(f"error: unsupported X-Format '{fmt}' (expected RGB565)", file=sys.stderr)
        return 1

    data = r.content
    expect = w * h * 2
    if len(data) != expect:
        print(f"error: body is {len(data)} bytes, expected {expect} ({w}x{h} RGB565)", file=sys.stderr)
        return 1

    dtype = ">u2" if args.byteswap else "<u2"
    px = np.frombuffer(data, dtype=dtype).reshape((h, w)).astype(np.uint32)
    r5 = (px >> 11) & 0x1F
    g6 = (px >> 5) & 0x3F
    b5 = px & 0x1F
    # Expand to 8-bit with bit replication so full white maps to 255.
    r8 = ((r5 << 3) | (r5 >> 2)).astype(np.uint8)
    g8 = ((g6 << 2) | (g6 >> 4)).astype(np.uint8)
    b8 = ((b5 << 3) | (b5 >> 2)).astype(np.uint8)
    if args.swap_rb:
        r8, b8 = b8, r8
    rgb = np.dstack([r8, g8, b8])
    Image.fromarray(rgb, "RGB").save(args.out)
    print(f"saved {args.out} ({w}x{h})")
    return 0


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