#!/usr/bin/env python3
"""Capture K1 LCD screen dump from J-Link RTT channel 1 and save as PNG."""

from __future__ import annotations

import argparse
import os
import struct
import sys
import time

MAGIC = b"SCRN"
HEADER_FMT = "<4sHHHH"
HEADER_SIZE = struct.calcsize(HEADER_FMT)
DEFAULT_DEVICES = ("AT32F403AC", "AT32F403A", "Cortex-M4")


def rgb565_bytes_to_rgb(hi: int, lo: int) -> tuple[int, int, int]:
    value = (hi << 8) | lo
    red = ((value >> 11) & 0x1F) << 3
    green = ((value >> 5) & 0x3F) << 2
    blue = (value & 0x1F) << 3
    return red, green, blue


def import_deps():
    try:
        import pylink  # noqa: F401
    except ImportError as exc:
        raise SystemExit("Missing dependency: pip install pylink-square") from exc
    try:
        from PIL import Image  # noqa: F401
    except ImportError as exc:
        raise SystemExit("Missing dependency: pip install pillow") from exc


def connect_jlink(device: str):
    import pylink

    jlink = pylink.JLink()
    jlink.open()
    jlink.set_tif(pylink.enums.JLinkInterfaces.SWD)

    last_error = None
    for candidate in (device, *DEFAULT_DEVICES):
        try:
            jlink.connect(candidate)
            print(f"Connected as {candidate}")
            return jlink
        except pylink.errors.JLinkException as exc:
            last_error = exc
    jlink.close()
    raise SystemExit(f"Failed to connect J-Link: {last_error}")


def find_rtt_control_block(jlink, ram_base: int = 0x20000000, ram_size: int = 0x18000) -> int | None:
    # AT32F403ARCT7 physical SRAM is 96KB (0x18000).
    """Locate SEGGER RTT CB in SRAM (needed when auto-scan misses high BSS)."""
    needle = b"SEGGER RTT"
    chunk = 0x1000
    was_halted = jlink.halted()
    if not was_halted:
        jlink.halt()
    try:
        for off in range(0, ram_size, chunk):
            data = bytes(jlink.memory_read8(ram_base + off, min(chunk, ram_size - off)))
            idx = data.find(needle)
            if idx >= 0:
                return ram_base + off + idx
    finally:
        if not was_halted:
            # Resume only — never reset (would wipe the on-screen UI).
            if hasattr(jlink, "restart"):
                jlink.restart()
            else:
                jlink.go()
            time.sleep(0.1)
    return None


def wait_rtt_ready(jlink, timeout_s: float = 15.0) -> None:
    deadline = time.time() + timeout_s
    while time.time() < deadline:
        try:
            if jlink.rtt_get_num_up_buffers() > 0:
                return
        except Exception:
            pass
        time.sleep(0.2)
    raise SystemExit("RTT control block found but buffers not ready")


def capture_rtt_png(out_path: str, device: str, timeout_s: float, do_reset: bool = False) -> None:
    import pylink
    from PIL import Image

    jlink = connect_jlink(device)
    try:
        if do_reset:
            jlink.reset(halt=False)
            time.sleep(2.0)

        cb = find_rtt_control_block(jlink)
        if cb is None:
            raise SystemExit("SEGGER RTT control block not found in SRAM")
        print(f"RTT CB @ 0x{cb:08X}")

        jlink.rtt_start(cb)
        wait_rtt_ready(jlink)

        # Drain terminal (ch0) and show firmware ACK lines if any.
        for _ in range(10):
            term = jlink.rtt_read(0, 1024)
            if term:
                print("RTT0:", bytes(term).decode("utf-8", errors="replace").strip())
            time.sleep(0.05)

        print("RTT ready. Sending DUMP command...")

        wrote = 0
        for _ in range(30):
            wrote = jlink.rtt_write(0, list(b"DUMP\n"))
            if wrote > 0:
                break
            time.sleep(0.2)
        if wrote <= 0:
            raise SystemExit("Failed to send DUMP command on RTT down channel 0")
        print(f"Sent DUMP ({wrote} bytes). Capturing screen...")

        buffer = b""
        deadline = time.time() + timeout_s

        while time.time() < deadline:
            term = jlink.rtt_read(0, 1024)
            if term:
                print("RTT0:", bytes(term).decode("utf-8", errors="replace").strip())
            chunk = jlink.rtt_read(1, 8192)
            if chunk:
                buffer += bytes(chunk)
                if len(buffer) >= HEADER_SIZE and buffer[:4] == MAGIC:
                    break
            time.sleep(0.01)

        if len(buffer) < HEADER_SIZE or buffer[:4] != MAGIC:
            raise SystemExit("Timeout waiting for SCRN header on RTT channel 1")

        _, width, height, fmt, _ = struct.unpack_from(HEADER_FMT, buffer, 0)
        if fmt != 1:
            raise SystemExit(f"Unsupported pixel format: {fmt}")

        pixel_bytes = width * height * 2
        buffer = buffer[HEADER_SIZE:]
        print(f"Receiving {width}x{height} RGB565 ({pixel_bytes} bytes)...")

        while len(buffer) < pixel_bytes and time.time() < deadline:
            chunk = jlink.rtt_read(1, 16384)
            if chunk:
                buffer += bytes(chunk)
            else:
                time.sleep(0.001)

        if len(buffer) < pixel_bytes:
            raise SystemExit(f"Incomplete dump: got {len(buffer)} / {pixel_bytes} bytes")

        image = Image.new("RGB", (width, height))
        pixels = image.load()
        offset = 0
        for y in range(height):
            for x in range(width):
                hi = buffer[offset]
                lo = buffer[offset + 1]
                offset += 2
                pixels[x, y] = rgb565_bytes_to_rgb(hi, lo)

        image.save(out_path)
        print(f"Saved {out_path}")
    finally:
        try:
            jlink.rtt_stop()
        except pylink.errors.JLinkException:
            pass
        jlink.close()


def main() -> None:
    parser = argparse.ArgumentParser(description="Capture K1 LCD via J-Link RTT")
    parser.add_argument("--out", default="screen.png", help="Output PNG path")
    parser.add_argument("--device", default="AT32F403AC", help="J-Link device name")
    parser.add_argument("--timeout", type=float, default=60.0, help="Seconds to wait")
    parser.add_argument("--open", action="store_true", help="Open PNG after capture")
    parser.add_argument(
        "--reset",
        action="store_true",
        help="Reset MCU before capture (default: keep current UI)",
    )
    args = parser.parse_args()

    import_deps()
    capture_rtt_png(args.out, args.device, args.timeout, do_reset=args.reset)

    if args.open and os.name == "nt":
        os.startfile(os.path.abspath(args.out))


if __name__ == "__main__":
    main()