#!/usr/bin/env python3
"""Flash tools/out/ui0902_res.bin to W25Q128 @ 0x00100000 via RTT."""

from __future__ import annotations

import argparse
import os
import sys
import time

DEFAULT_DEVICES = ("Cortex-M4", "AT32F403AC", "AT32F403A")
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEFAULT_BIN = os.path.join(ROOT, "tools", "out", "ui0902_res.bin")


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 = 30.0) -> str:
    deadline = time.time() + timeout_s
    buf = ""
    while time.time() < deadline:
        buf += rtt_read_text(jlink, 0.15)
        if marker in buf:
            return buf
    raise SystemExit(f"Timeout waiting for {marker!r}\n--- RTT ---\n{buf[-800:]}")


def send_bytes(jlink, payload: bytes) -> None:
    off = 0
    while off < len(payload):
        n = jlink.rtt_write(0, list(payload[off : off + 128]))
        if n <= 0:
            time.sleep(0.01)
            continue
        off += n


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--bin", default=DEFAULT_BIN)
    parser.add_argument("--device", default="Cortex-M4")
    parser.add_argument("--no-reset", action="store_true",
                        help="Skip hardware reset (device must already be running)")
    args = parser.parse_args()

    if not os.path.isfile(args.bin):
        raise SystemExit(f"missing bin: {args.bin}")

    data = open(args.bin, "rb").read()
    size = len(data)
    print(f"bin={args.bin} size={size}")

    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:
            # Hardware reset clears any stuck flash_mode from a previous attempt
            print("hardware reset...")
            jlink.reset(halt=False)
            time.sleep(3.5)
            # RTT CB may move after reboot — re-find
            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)

        cmd = f"flash ui0902 {size}\n".encode("ascii")
        send_bytes(jlink, cmd)
        print("sent flash command, waiting GO...")
        log = wait_for(jlink, "FLASH_UI0902_GO", 20.0)
        print(log.strip().splitlines()[-1])
        time.sleep(0.2)

        # Smaller chunks + brief pause so erase/write can keep up with RTT down-buffer
        chunk = 128
        sent = 0
        t0 = time.time()
        send_bytes(jlink, data[0:256])
        sent = 256
        wait_for(jlink, "FLASH_RX", 20.0)
        print("device receiving...")
        while sent < size:
            end = min(sent + chunk, size)
            send_bytes(jlink, data[sent:end])
            sent = end
            if sent % 256 == 0:
                time.sleep(0.002)
            if sent % 4096 == 0 or sent == size:
                log = wait_for(jlink, "FLASH_ACK", 60.0)
                if "FLASH_UI0902_OK" in log:
                    print(log.strip().splitlines()[-1])
                    print("Done. Mode-select should redraw with new assets.")
                    return
            if sent % (64 * 1024) == 0 or sent == size:
                elapsed = time.time() - t0
                print(f"  {sent}/{size} ({100.0 * sent / size:.1f}%) {elapsed:.1f}s")

        log = wait_for(jlink, "FLASH_UI0902_OK", 60.0)
        print(log.strip().splitlines()[-1])
        print("Done. Mode-select should redraw with new assets.")
    finally:
        try:
            jlink.rtt_stop()
        except Exception:
            pass
        jlink.close()


if __name__ == "__main__":
    main()