#!/usr/bin/env python3
"""Erase entire W25Q128 external Flash via RTT command: flash erase chip

Requires firmware that implements the RTT command (APP/app_log.c).
Chip erase typically takes 30–120s; do not power-cycle mid-erase.
"""

from __future__ import annotations

import argparse
import time

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


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 Exception 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):
    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 rtt_read_text(jlink, timeout_s: float = 0.2) -> str:
    deadline = time.time() + timeout_s
    chunks: list[bytes] = []
    while time.time() < deadline:
        try:
            data = bytes(jlink.rtt_read(0, 512) or [])
        except Exception:
            data = b""
        if data:
            chunks.append(data)
            deadline = time.time() + timeout_s
        else:
            time.sleep(0.02)
    return b"".join(chunks).decode("ascii", errors="replace")


def wait_for(jlink, marker: str, timeout_s: float) -> str:
    deadline = time.time() + timeout_s
    buf = ""
    last_print = 0.0
    while time.time() < deadline:
        chunk = rtt_read_text(jlink, 0.2)
        if chunk:
            buf += chunk
            if marker in buf:
                return buf
        now = time.time()
        if now - last_print >= 5.0:
            elapsed = now - (deadline - timeout_s)
            print(f"  ... waiting erase ({elapsed:.0f}s)")
            last_print = now
        time.sleep(0.05)
    raise SystemExit(f"Timeout waiting for {marker!r}\n--- RTT ---\n{buf[-1200:]}")


def send_bytes(jlink, payload: bytes) -> None:
    off = 0
    while off < len(payload):
        try:
            n = jlink.rtt_write(0, list(payload[off : off + 64]))
        except Exception:
            time.sleep(0.05)
            continue
        if n is None or n <= 0:
            time.sleep(0.02)
            continue
        off += n


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--device", default="Cortex-M4")
    parser.add_argument("--no-reset", action="store_true")
    parser.add_argument("--timeout", type=float, default=180.0,
                        help="seconds to wait for FLASH_ERASE_OK (default 180)")
    args = parser.parse_args()

    jlink = connect_jlink(args.device)
    try:
        cb = find_rtt_control_block(jlink)
        if cb is None:
            raise SystemExit("RTT control block not found")
        print(f"RTT CB @ 0x{cb:08X}")
        jlink.rtt_start(cb)
        time.sleep(0.3)
        _ = rtt_read_text(jlink, 0.3)

        if not args.no_reset:
            print("hardware reset...")
            jlink.reset(halt=False)
            time.sleep(3.5)
            try:
                jlink.rtt_stop()
            except Exception:
                pass
            cb = find_rtt_control_block(jlink)
            if cb is None:
                raise SystemExit("RTT CB missing after reset")
            print(f"RTT CB @ 0x{cb:08X}")
            jlink.rtt_start(cb)
            time.sleep(0.5)
            _ = rtt_read_text(jlink, 0.5)

        print("Sending: flash erase chip")
        send_bytes(jlink, b"flash erase chip\n")
        log = wait_for(jlink, "FLASH_ERASE_BEGIN", 20.0)
        print(log.strip().splitlines()[-1])
        print(f"Chip erase running (timeout {args.timeout:.0f}s)...")
        log = wait_for(jlink, "FLASH_ERASE_OK", args.timeout)
        print(log.strip().splitlines()[-1])
        print("ExtFlash chip erase done. Flash MCU + extflash_ALL*.res manually.")
    finally:
        try:
            jlink.rtt_stop()
        except Exception:
            pass
        jlink.close()


if __name__ == "__main__":
    main()