#!/usr/bin/env python3
"""Set up a Seeed SenseCAP T1000-E with Stolo Node firmware.

One command does everything a fresh tracker needs:

    python3 stolo_t1000e_setup.py --flash stolo-t1000e-firmware.zip

  1. flashes the firmware over USB (needs `pip install adafruit-nrfutil`),
  2. writes the radio's identity (product 0x1E, model 0xB5, 863-928 MHz)
     if it has none yet,
  3. records the firmware's checksum, which the radio requires before it
     will switch its LoRa transmitter on,
  4. turns Bluetooth on, so a phone can pair.

Run it again without --flash to repair steps 2-4. Requires pyserial
(`pip install pyserial`). Close any app that has the serial port open, and
switch off Bluetooth on paired phones while it runs: the firmware ignores
USB while a phone is connected over Bluetooth.

GPL-3.0-or-later. Part of the Stolo Node firmware release for the T1000-E.
"""
import argparse, hashlib, os, struct, subprocess, sys, time, zipfile

try:
    import serial
    from serial.tools import list_ports
except ImportError:
    sys.exit("pyserial is required: pip install pyserial")

FEND, FESC, TFEND, TFESC = 0xC0, 0xDB, 0xDC, 0xDD
CMD_DETECT, DETECT_REQ, DETECT_RESP = 0x08, 0x73, 0x46
CMD_ROM_READ, CMD_ROM_WRITE = 0x51, 0x52
CMD_FW_HASH, CMD_HASHES, CMD_BT_CTRL = 0x58, 0x60, 0x46
CMD_RESET, RESET_BYTE = 0x55, 0xF8
ADDR_PRODUCT, ADDR_MODEL, ADDR_HW_REV, ADDR_SERIAL, ADDR_MADE = 0x00, 0x01, 0x02, 0x03, 0x07
ADDR_CHKSUM, ADDR_INFO_LOCK, ADDR_CONF_BT = 0x0B, 0x9B, 0xB0
LOCK_BYTE = BT_ENABLE = 0x73
PRODUCT_T1000E, MODEL_B5, HW_REV = 0x1E, 0xB5, 0x01
SEEED_VIDS = (0x2886, 0x239A)
# The nRF52 EEPROM is emulated in flash and stores one byte at a time;
# writes sent faster than this are dropped.
WRITE_PACING_S = 0.06


def say(msg):
    print(msg, flush=True)


def esc(data):
    out = bytearray()
    for b in data:
        if b == FEND: out += bytes([FESC, TFEND])
        elif b == FESC: out += bytes([FESC, TFESC])
        else: out.append(b)
    return bytes(out)


def frames(raw):
    """Unescaped KISS frames (command byte first) in a byte string."""
    out = []
    for chunk in raw.split(bytes([FEND])):
        if not chunk: continue
        f, e = bytearray(), False
        for b in chunk:
            if e: f.append(FEND if b == TFEND else FESC if b == TFESC else b); e = False
            elif b == FESC: e = True
            else: f.append(b)
        out.append(bytes(f))
    return out


def find_port(explicit):
    if explicit: return explicit
    for p in list_ports.comports():
        if p.vid in SEEED_VIDS: return p.device
    sys.exit("No T1000-E found on USB. Plug it in (and switch it on), or pass --port.")


def wait_for_port(port, timeout=30):
    t = time.time()
    while time.time() - t < timeout:
        try:
            # write_timeout: a radio that stopped reading USB must not hang us.
            return serial.Serial(port, 115200, timeout=0.2, write_timeout=2)
        except (serial.SerialException, OSError):
            time.sleep(0.2)
    return None


class Radio:
    def __init__(self, port):
        self.port = port
        self.ser = wait_for_port(port)
        if self.ser is None: sys.exit(f"Could not open {port}.")

    def close(self):
        try: self.ser.close()
        except Exception: pass

    def send(self, cmd, payload=b""):
        self.ser.write(bytes([FEND, cmd]) + esc(payload) + bytes([FEND])); self.ser.flush()

    def collect(self, seconds):
        d, t = b"", time.time()
        while time.time() - t < seconds:
            try: d += self.ser.read(512)
            except (serial.SerialException, OSError): break
        return frames(d)

    def detect(self, timeout=25):
        t = time.time()
        while time.time() - t < timeout:
            try:
                self.ser.reset_input_buffer()
                self.send(CMD_DETECT, bytes([DETECT_REQ]))
                if any(f[:2] == bytes([CMD_DETECT, DETECT_RESP]) for f in self.collect(0.4)):
                    return True
            except serial.SerialTimeoutException:
                time.sleep(0.5)   # not reading USB yet (booting, or a phone holds it)
            except (serial.SerialException, OSError):
                self.close(); time.sleep(0.5); self.ser = wait_for_port(self.port) or self.ser
        return False

    def read_eeprom(self):
        for _ in range(4):
            self.ser.reset_input_buffer()
            self.send(CMD_ROM_READ, b"\x00")
            for f in self.collect(1.5):
                if f[:1] == bytes([CMD_ROM_READ]) and len(f) > ADDR_CONF_BT:
                    return f[1:]
        return None

    def write_byte(self, addr, value):
        self.send(CMD_ROM_WRITE, bytes([addr, value])); time.sleep(WRITE_PACING_S)

    def hashes(self):
        self.ser.reset_input_buffer()
        self.send(CMD_HASHES, b"\x01"); time.sleep(0.2); self.send(CMD_HASHES, b"\x02")
        got = {}
        for f in self.collect(1.5):
            if f[:1] == bytes([CMD_HASHES]) and len(f) >= 34: got[f[1]] = f[2:34]
        return got.get(1), got.get(2)


def identity_bytes():
    serial_no = os.urandom(4)
    made = struct.pack(">I", int(time.time()))
    info = bytes([PRODUCT_T1000E, MODEL_B5, HW_REV]) + serial_no + made
    return info, hashlib.md5(info).digest()


def eeprom_ok(rom):
    if rom[ADDR_INFO_LOCK] != LOCK_BYTE: return False
    if rom[ADDR_PRODUCT] != PRODUCT_T1000E or rom[ADDR_MODEL] != MODEL_B5: return False
    return hashlib.md5(bytes(rom[0:11])).digest() == bytes(rom[ADDR_CHKSUM:ADDR_CHKSUM + 16])


def provision(radio):
    rom = radio.read_eeprom()
    if rom is None:
        sys.exit("The radio did not return its EEPROM. Is Stolo Node firmware installed? Is a phone connected over Bluetooth?")
    if eeprom_ok(rom):
        say("Identity: already set (SenseCAP T1000-E, 863-928 MHz).")
        return
    if rom[ADDR_INFO_LOCK] == LOCK_BYTE:
        sys.exit("The radio has a locked identity that is not a T1000-E identity. Not changing it.")
    info, checksum = identity_bytes()
    plan = {ADDR_PRODUCT + i: b for i, b in enumerate(info)}
    plan.update({ADDR_CHKSUM + i: b for i, b in enumerate(checksum)})
    say("Identity: writing (SenseCAP T1000-E, 863-928 MHz)...")
    for attempt in range(4):
        pending = {a: v for a, v in plan.items() if rom is None or rom[a] != v}
        if not pending: break
        for a, v in sorted(pending.items()): radio.write_byte(a, v)
        time.sleep(0.5); rom = radio.read_eeprom()
    else:
        sys.exit("Could not write the identity reliably. Run the tool again.")
    for attempt in range(4):
        radio.write_byte(ADDR_INFO_LOCK, LOCK_BYTE); time.sleep(0.6)
        rom = radio.read_eeprom()
        if rom and eeprom_ok(rom): say("Identity: written and verified."); return
    sys.exit("The identity was written but could not be locked. Run the tool again.")


def firmware_hash_from_zip(path):
    with zipfile.ZipFile(path) as z:
        bins = [n for n in z.namelist() if n.endswith(".bin")]
        if len(bins) != 1: sys.exit(f"{path} does not look like a T1000-E firmware package.")
        return hashlib.sha256(z.read(bins[0])).digest()


def flash(port, package):
    say("Flashing (the radio restarts into its bootloader; this takes about 20 s)...")
    r = subprocess.run(["adafruit-nrfutil", "dfu", "serial", "--package", package, "-p", port,
                        "-b", "115200", "--singlebank", "--touch", "1200"], capture_output=True, text=True)
    if r.returncode != 0:
        sys.exit("Flashing failed:\n" + (r.stdout + r.stderr)[-800:] +
                 "\nIf adafruit-nrfutil is missing: pip install adafruit-nrfutil")
    say("Flashed.")


def main():
    ap = argparse.ArgumentParser(description="Set up a SenseCAP T1000-E with Stolo Node firmware.")
    ap.add_argument("--port", help="serial port (found automatically when only one T1000-E is plugged in)")
    ap.add_argument("--flash", metavar="ZIP", help="firmware package to flash first")
    ap.add_argument("--firmware", metavar="ZIP", help="package the radio is running (for the checksum) when not flashing")
    a = ap.parse_args()
    package = a.flash or a.firmware
    if not package: ap.error("pass --flash ZIP (or --firmware ZIP for a radio that is already flashed)")
    if not os.path.isfile(package): sys.exit(f"{package}: no such file")
    target_hash = firmware_hash_from_zip(package)
    port = find_port(a.port)
    if a.flash:
        flash(port, a.flash); time.sleep(3); port = find_port(a.port)

    radio = Radio(port)
    if not radio.detect():
        sys.exit("The radio does not answer on USB. If a phone is connected to it over Bluetooth, "
                 "switch that phone's Bluetooth off and run this again.")
    provision(radio)

    # Bluetooth first: the radio checks its firmware at boot, and the
    # checksum restart below is the boot that has to see Bluetooth on.
    rom = radio.read_eeprom()
    if rom is not None and rom[ADDR_CONF_BT] == BT_ENABLE:
        say("Bluetooth: on.")
    else:
        radio.send(CMD_BT_CTRL, b"\x01"); time.sleep(1.0)
        rom = radio.read_eeprom()
        if rom is None or rom[ADDR_CONF_BT] != BT_ENABLE:
            sys.exit("Could not switch Bluetooth on. Run the tool again.")
        say("Bluetooth: on.")

    def restart(r):
        r.close(); time.sleep(3)
        r2 = Radio(find_port(a.port))
        if not r2.detect(): sys.exit("The radio did not come back after restarting.")
        return r2

    target, live = radio.hashes()
    if target != target_hash:
        say("Firmware checksum: recording (the radio restarts)...")
        radio.send(CMD_FW_HASH, target_hash); time.sleep(1.0)
        radio = restart(radio)
        target, live = radio.hashes()
        if target != target_hash: sys.exit("The checksum did not stick. Run the tool again.")
    # What the radio computed at boot must match, or it will not transmit.
    if live != target_hash:
        radio.send(CMD_RESET, bytes([RESET_BYTE])); time.sleep(0.5)
        radio = restart(radio)
        target, live = radio.hashes()
    if live != target_hash:
        sys.exit("The radio's firmware check does not match this package. Was it flashed with a different file?")
    say("Firmware check: passed (the radio may transmit).")
    radio.close()
    say("\nDone. Unplug the radio, hold its button for 1 s to switch it on, then pair it from the Stolo app\n"
        "(it accepts a new phone for 2 minutes after power-on; press the button twice to reopen that window).")


if __name__ == "__main__":
    try:
        main()
    except serial.SerialTimeoutException:
        sys.exit("The radio stopped reading USB. Unplug it, switch off Bluetooth on any paired phone, "
                 "plug it back in and run this again.")
    except KeyboardInterrupt:
        sys.exit(1)