#!/usr/bin/env python3
"""Send sys reset via J-Link RTT to reboot K1 MCU (software NVIC reset)."""

from __future__ import annotations

import argparse
import time

DEFAULT_DEVICES = ("Cortex-M4", "AT32F403AC", "AT32F403A")
RESET_CMD = b"sys reset\n"
ACK = b"SYS_RESET_OK"


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:
            resume_cpu(jlink)
    return None


def resume_cpu(jlink) -> None:
    """J-Link often halts the core on reset; must run() or UI stays black."""
    try:
        jlink.rtt_stop()
    except Exception:
        pass
    try:
        jlink.reset(halt=False)
    except Exception:
        pass
    try:
        if hasattr(jlink, "restart"):
            jlink.restart()
        else:
            jlink.go()
    except Exception:
        jlink.go()
    time.sleep(0.05)


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_reset(device: str, timeout_s: float, wait_boot_s: float, hw_fallback: bool) -> None:
    jlink = connect_jlink(device)
    ack_seen = False
    try:
        if jlink.halted():
            print("Target was halted, resuming before RTT...")
            resume_cpu(jlink)
            time.sleep(0.2)

        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}")

        resume_cpu(jlink)
        time.sleep(0.15)
        jlink.rtt_start(cb)
        wait_rtt_ready(jlink)

        wrote = 0
        for _ in range(30):
            wrote = jlink.rtt_write(0, list(RESET_CMD))
            if wrote > 0:
                break
            time.sleep(0.2)
        if wrote <= 0:
            raise SystemExit("Failed to send sys reset on RTT down channel 0")
        print(f"Sent sys reset ({wrote} bytes)")

        deadline = time.time() + timeout_s
        while time.time() < deadline:
            chunk = jlink.rtt_read(0, 1024)
            if chunk:
                text = bytes(chunk)
                if ACK in text:
                    ack_seen = True
                    print("SYS_RESET_OK (firmware handled reset)")
                    break
            time.sleep(0.05)

        if not ack_seen:
            print("No SYS_RESET_OK from firmware.")
            if hw_fallback:
                print("Using J-Link hardware reset fallback...")
                resume_cpu(jlink)
            else:
                print("Hint: re-run with --hw-fallback or flash latest firmware.")

        # Critical: release CPU after reset — otherwise screen stays black.
        print("Resuming CPU after reset...")
        resume_cpu(jlink)

        if wait_boot_s > 0:
            print(f"Waiting {wait_boot_s:.0f}s for auto boot...")
            time.sleep(wait_boot_s)
            resume_cpu(jlink)
            print("Done. Device should be back on mode-select screen.")
    finally:
        try:
            resume_cpu(jlink)
        except Exception:
            pass
        jlink.close()


def main() -> None:
    parser = argparse.ArgumentParser(description="Reboot K1 via RTT sys reset command")
    parser.add_argument("--device", default="AT32F403AC", help="J-Link device name")
    parser.add_argument("--timeout", type=float, default=3.0, help="Seconds to wait for ACK")
    parser.add_argument(
        "--wait-boot",
        type=float,
        default=4.0,
        help="Seconds to wait after reset for auto power-on (0 to skip)",
    )
    parser.add_argument(
        "--no-hw-fallback",
        action="store_true",
        help="Do not use J-Link hardware reset when firmware ACK is missing",
    )
    args = parser.parse_args()

    import_deps()
    send_reset(args.device, args.timeout, args.wait_boot, not args.no_hw_fallback)


if __name__ == "__main__":
    main()