#!/usr/bin/env python3
"""Robust RT-4D application flasher (bootloader / 0x39 protocol).

Replicates the official DualTachyon flasher-cli handshake:
  1. 0xFF probe  — spam 0xFF until the bootloader echoes 0xFF
  2. erase       — two-part official-updater sequence (checksum seed 0x48)
                     [0x39,0x33,0x05,0x10] then [0x39,0x33,0x05,0x55]
  3. write       — [0x57, off>>8, off] + 1024B, checksum seed 0x48, ACK 0x06

Only the APPLICATION region is erased/written; the bootloader at
0x08000000..0x08002800 is never touched, so a failed flash is always
recoverable by reflashing the stock app image.

Usage: python flash_rt4d.py COM7 <firmware.bin>
"""
import sys, time, serial

SEED = 0x48
ACK = 0x06


def cksum(payload):
    s = SEED
    for b in payload:
        s = (s + b) & 0xFF
    return bytes(payload) + bytes([s])


def handshake(p, timeout=5.0):
    p.reset_input_buffer()
    t0 = time.time()
    while time.time() - t0 < timeout:
        b = p.read(1)
        if not b:
            p.write(bytes([0xFF]))
        elif b == b"\xff":
            return True
        # ignore any other stray byte and keep probing
    return False


def erase(p):
    for trig in (0x10, 0x55):
        p.reset_input_buffer()
        p.write(cksum([0x39, 0x33, 0x05, trig]))
        # wait up to ~4s for ACK
        ok = False
        t0 = time.time()
        while time.time() - t0 < 4.0:
            b = p.read(1)
            if b == bytes([ACK]):
                ok = True
                break
            if b:  # unexpected byte
                pass
        if not ok:
            raise SystemExit(f"[E] erase phase 0x{trig:02X}: no ACK")
        print(f"[i] erase 0x{trig:02X}: ACK")


def write_fw(p, fw):
    # pad up to a 1024 multiple with 0xFF (erased-flash value)
    pad = (-len(fw)) % 1024
    data = fw + b"\xff" * pad
    total = len(data)
    for off in range(0, total, 1024):
        chunk = data[off:off + 1024]
        payload = bytes([0x57, (off >> 8) & 0xFF, off & 0xFF]) + chunk
        p.reset_input_buffer()
        p.write(cksum(payload))
        # ACK
        ok = False
        t0 = time.time()
        while time.time() - t0 < 2.0:
            b = p.read(1)
            if b == bytes([ACK]):
                ok = True
                break
        if not ok:
            raise SystemExit(f"\n[E] write at 0x{off:05X}: no ACK "
                             f"(wrote {off}/{total})")
        pct = (off + len(chunk)) * 100 // total
        print(f"\r[i] writing 0x{off:05X}/{total:#07x} ({pct:3d}%)", end="", flush=True)
    print("\n[i] write complete")


def main():
    if len(sys.argv) != 3:
        raise SystemExit(__doc__)
    port, path = sys.argv[1], sys.argv[2]
    fw = open(path, "rb").read()
    print(f"[i] firmware {path}: {len(fw)} bytes (0x{len(fw):X})")
    p = serial.Serial(port, 115200, timeout=0.3, write_timeout=3)
    try:
        print("[i] handshake (0xFF probe)...")
        if not handshake(p):
            raise SystemExit("[E] no bootloader (radio not in flash mode). "
                             "Power off, hold * (or the flash-mode key) and power on.")
        print("[i] bootloader ready")
        print("[i] erasing application region...")
        erase(p)
        print("[i] flashing...")
        write_fw(p, fw)
        print("[i] DONE. Power-cycle the radio normally.")
    finally:
        p.close()


if __name__ == "__main__":
    main()