#!/usr/bin/env python3
"""Recover K1 from J-Link halt / black screen: hardware reset + run (no RTT required)."""

from __future__ import annotations

import argparse
import time

DEFAULT_DEVICES = ("Cortex-M4", "AT32F403AC", "AT32F403A")


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 run_go(jlink) -> None:
    try:
        jlink.reset(halt=False)
    except Exception:
        pass
    try:
        if hasattr(jlink, "restart"):
            jlink.restart()
        else:
            jlink.go()
    except Exception:
        jlink.go()


def recover(device: str, wait_boot_s: float) -> None:
    jlink = connect_jlink(device)
    try:
        halted = jlink.halted()
        print(f"Target halted={halted}")
        print("Reset + run...")
        run_go(jlink)
        if wait_boot_s > 0:
            print(f"Waiting {wait_boot_s:.0f}s for auto boot...")
            time.sleep(wait_boot_s)
        run_go(jlink)
        print("Done. Screen should show UI; if still black, run flash_app.bat")
    finally:
        try:
            run_go(jlink)
        except Exception:
            pass
        jlink.close()


def main() -> None:
    parser = argparse.ArgumentParser(description="Recover K1 from halt/black screen via J-Link")
    parser.add_argument("--device", default="AT32F403AC")
    parser.add_argument("--wait-boot", type=float, default=4.0)
    args = parser.parse_args()
    import_deps()
    recover(args.device, args.wait_boot)


if __name__ == "__main__":
    main()