#!/usr/bin/env python3
"""Pull K1 device log ring buffer via J-Link RTT (log dump command)."""

from __future__ import annotations

import argparse
import os
import sys
import time
from datetime import datetime

DEFAULT_DEVICES = ("Cortex-M4", "AT32F403AC", "AT32F403A")
DUMP_CMD = b"log dump\n"
PROBE_CMD = b"tp status\n"
END_MARKERS = (b"LOG_DUMP_END",)


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}")
            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 = 30) -> None:
    for _ in range(retries):
        wrote = jlink.rtt_write(0, list(payload))
        if wrote > 0:
            print(f"Sent command ({wrote} bytes): {payload.decode('ascii', errors='replace').strip()}")
            return
        time.sleep(0.2)
    raise SystemExit("Failed to send RTT down command on channel 0")


def dump_log(out_path: str, device: str, timeout_s: float) -> None:
    jlink = connect_jlink(device)
    lines: list[str] = []
    try:
        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 stale RTT output.
        for _ in range(5):
            stale = jlink.rtt_read(0, 4096)
            if stale:
                text = bytes(stale).decode("utf-8", errors="replace")
                if text.strip():
                    print("RTT0 stale:", text.strip()[:200])
            time.sleep(0.05)

        send_command(jlink, PROBE_CMD)
        time.sleep(0.25)
        send_command(jlink, DUMP_CMD)

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

        while time.time() < deadline:
            chunk = jlink.rtt_read(0, 4096)
            if chunk:
                buffer += bytes(chunk)
                if not started and b"LOG_DUMP_BEGIN" in buffer:
                    started = True
                if b"LOG_DUMP_END" in buffer:
                    break
            else:
                time.sleep(0.01)

        if b"LOG_DUMP_END" not in buffer:
            raise SystemExit("Timeout waiting for LOG_DUMP_END (is firmware with app_log flashed?)")

        text = buffer.decode("utf-8", errors="replace")
        capture = False
        for line in text.splitlines():
            if line.strip() == "LOG_DUMP_BEGIN":
                capture = True
                continue
            if line.strip() == "LOG_DUMP_END":
                break
            if capture:
                lines.append(line)

        header = (
            f"# K1 log dump {datetime.now().isoformat(timespec='seconds')}\n"
            f"# device={device}\n"
        )
        body = "\n".join(lines) + ("\n" if lines else "")
        with open(out_path, "w", encoding="utf-8", newline="\n") as f:
            f.write(header)
            f.write(body)

        print(f"Saved {len(lines)} log lines to {out_path}")
    finally:
        try:
            jlink.rtt_stop()
        except Exception:
            pass
        jlink.close()


def main() -> None:
    parser = argparse.ArgumentParser(description="Dump K1 RAM log ring via J-Link RTT")
    parser.add_argument(
        "--out",
        default=f"k1_log_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log",
        help="Output log file path",
    )
    parser.add_argument("--device", default="AT32F403AC", help="J-Link device name")
    parser.add_argument("--timeout", type=float, default=20.0, help="Seconds to wait for dump")
    args = parser.parse_args()

    import_deps()
    dump_log(args.out, args.device, args.timeout)

    if os.name == "nt" and os.path.isfile(args.out):
        print(f"Full path: {os.path.abspath(args.out)}")


if __name__ == "__main__":
    main()