#!/usr/bin/env python3
"""Live-stream K1 touch logs via J-Link RTT while you touch the screen.

Usage:
  python rtt_touch_live.py
  python rtt_touch_live.py --all          # show every RTT line
  python rtt_touch_live.py --seconds 120  # auto-stop after N seconds
"""

from __future__ import annotations

import argparse
import sys
import time

DEFAULT_DEVICES = ("Cortex-M4", "AT32F403AC", "AT32F403A")
LIVE_ON_CMD = b"tp live\n"
LIVE_OFF_CMD = b"tp live off\n"

# Default: touch / UI hit / message path
DEFAULT_KEYS = ("[TP", "[UI", "[MSG")


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


def connect_jlink(device: str):
    import pylink

    jlink = pylink.JLink()
    jlink.open()
    try:
        jlink.exec_command("HideDeviceSelection = 1")
    except Exception:
        pass
    jlink.set_tif(pylink.enums.JLinkInterfaces.SWD)

    last_error = None
    for candidate in (device, *DEFAULT_DEVICES):
        try:
            try:
                jlink.exec_command(f"Device = {candidate}")
            except Exception:
                pass
            jlink.connect(candidate)
            print(f"Connected as {candidate}", flush=True)
            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:
    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:
            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 send_command(jlink, payload: bytes, retries: int = 40) -> None:
    for _ in range(retries):
        wrote = jlink.rtt_write(0, list(payload))
        if wrote > 0:
            print(f">> {payload.decode('ascii', errors='replace').strip()}", flush=True)
            return
        time.sleep(0.15)
    print("WARN: failed to send RTT command", flush=True)


def want_line(line: str, show_all: bool, keys: tuple[str, ...]) -> bool:
    if show_all:
        return True
    return any(k in line for k in keys)


def run_live(device: str, show_all: bool, seconds: float, out_path: str | None) -> None:
    jlink = connect_jlink(device)
    out_fp = open(out_path, "w", encoding="utf-8", newline="\n") if out_path else None
    try:
        if jlink.halted():
            if hasattr(jlink, "restart"):
                jlink.restart()
            else:
                jlink.go()
            time.sleep(0.3)

        cb = find_rtt_control_block(jlink)
        if cb is None:
            raise SystemExit("SEGGER RTT control block not found — is app running?")
        print(f"RTT CB @ 0x{cb:08X}", flush=True)

        jlink.rtt_start(cb)
        wait_rtt_ready(jlink)

        # Drain backlog
        for _ in range(8):
            junk = jlink.rtt_read(0, 4096)
            if not junk:
                break

        print("Waiting ~8s for UI/touch task...", flush=True)
        time.sleep(8.0)
        for _ in range(4):
            junk = jlink.rtt_read(0, 4096)
            if not junk:
                break

        # Event-only: firmware no longer floods idle samples.
        # Optional: tp live just tags the session; leave it off by default.
        print("Event log ON — only touch events will print. Touch the screen. Ctrl+C to stop.", flush=True)
        print("-" * 60, flush=True)

        buf = b""
        t0 = time.time()
        while True:
            if seconds > 0 and (time.time() - t0) >= seconds:
                print("\n(time limit reached)", flush=True)
                break

            chunk = jlink.rtt_read(0, 4096)
            if chunk:
                buf += bytes(chunk)
                while True:
                    nl = buf.find(b"\n")
                    if nl < 0:
                        break
                    raw = buf[:nl]
                    buf = buf[nl + 1 :]
                    line = raw.decode("utf-8", errors="replace").rstrip("\r")
                    if not line:
                        continue
                    if want_line(line, show_all, DEFAULT_KEYS):
                        print(line, flush=True)
                        if out_fp:
                            out_fp.write(line + "\n")
                            out_fp.flush()
            else:
                time.sleep(0.01)
    except KeyboardInterrupt:
        print("\nStopped by user.", flush=True)
    finally:
        try:
            jlink.rtt_stop()
        except Exception:
            pass
        try:
            if hasattr(jlink, "restart"):
                jlink.restart()
            else:
                jlink.go()
        except Exception:
            pass
        jlink.close()
        if out_fp:
            out_fp.close()
            print(f"Saved: {out_path}", flush=True)


def main() -> None:
    parser = argparse.ArgumentParser(description="Live K1 touch RTT monitor")
    parser.add_argument("--device", default="AT32F403AC")
    parser.add_argument("--all", action="store_true", help="Print all RTT lines")
    parser.add_argument("--seconds", type=float, default=0, help="Auto stop after N seconds (0=forever)")
    parser.add_argument(
        "--out",
        default="",
        help="Also append filtered lines to this file",
    )
    args = parser.parse_args()
    import_deps()
    run_live(args.device, args.all, args.seconds, args.out or None)


if __name__ == "__main__":
    main()