#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
ble_ota_update.py - PC-side OTA updater for CAIIC-MCM devices, transport
protocol per docs/ble_protocol.md section 6.6 (0xCA frames, OTA_BEGIN/DATA/
END/ABORT -> framed OTA_RSP acks).

Two channels:
    BLE (default):  dedicated OTA characteristic ...e0005 - one frame per
                    write-with-response, the framed OTA_RSP is read back from
                    the same characteristic (lockstep, no notify needed).
    UART (--uart):  e.g. --uart COM4 - the CLI command "ota" switches the
                    serial link to binary frame mode (marker-confirmed
                    handshake; OTA_ABORT exits); same frames on the wire.

Both channels select the payload by CUR_BANK (info item 0x07, read-only
characteristic ...e0004 / CLI "devinfo"): OTA always writes the INACTIVE
bank with the image linked for it, from the single-file package
(mothercup_ble_ota.bin, 52B header + both payloads).

    usage:  ble_ota.bat [mothercup_ble_ota.bin]
            ble_ota.bat --uart COM4 [mothercup_ble_ota.bin]

Success = the device resets after OTA_END and comes back on the new bank
with the new version (verified by a second CUR_BANK/FW_VERSION read).
"""

import asyncio
import os
import re
import struct
import sys
import time
import zlib

NAME_PREFIX = "CAIIC-MCM-20260902"
OTA_RW_UUID = "00002760-08c2-11e1-9073-0e8ac72e0005"   # OTA: write frame, read RSP
INFO_RD_UUID = "00002760-08c2-11e1-9073-0e8ac72e0004"  # Read-only info TLV

TYPE_OTA_BEGIN = 0x10
TYPE_OTA_DATA = 0x11
TYPE_OTA_END = 0x12
TYPE_OTA_ABORT = 0x13
TYPE_OTA_RSP = 0x1F

COMBO_MAGIC = 0xCA10BA11
COMBO_HDR_LEN = 52

RSP_TIMEOUT_S = 2.0     # one lockstep round trip (flash erase adds ~50ms)
XFER_TRIES    = 3       # resend a frame whose ack was lost (flaky links)
REBOOT_WAIT_S = 20.0

ST_OK        = 0
ST_BAD_STATE = 2

_seq = [0]


def crc32(data):
    return zlib.crc32(data) & 0xFFFFFFFF


def crc16_ccitt(data):
    crc = 0xFFFF
    for b in data:
        crc ^= b << 8
        for _ in range(8):
            crc = ((crc << 1) ^ 0x1021) & 0xFFFF if crc & 0x8000 else (crc << 1) & 0xFFFF
    return crc


def encode_frame(ftype, payload):
    out = bytearray([0xCA, ftype & 0xFF, _seq[0] & 0xFF,
                     len(payload) & 0xFF, (len(payload) >> 8) & 0xFF])
    out += bytes(payload)
    crc = crc16_ccitt(out[1:])
    out += bytes([crc & 0xFF, (crc >> 8) & 0xFF])
    _seq[0] = (_seq[0] + 1) & 0xFF
    return bytes(out)


class FrameDecoder:
    """0xCA byte-stream reassembly (same rules as the firmware)."""

    def __init__(self):
        self.buf = bytearray()

    def feed(self, data):
        frames = []
        self.buf += bytes(data)
        while True:
            while self.buf and self.buf[0] != 0xCA:
                del self.buf[0]
            if len(self.buf) < 5:
                break
            plen = self.buf[3] | (self.buf[4] << 8)
            if plen > 480:
                del self.buf[0]
                continue
            total = 5 + plen + 2
            if len(self.buf) < total:
                break
            crc = crc16_ccitt(self.buf[1:5 + plen])
            rx = self.buf[5 + plen] | (self.buf[5 + plen + 1] << 8)
            if crc != rx:
                del self.buf[0]
                continue
            frames.append((self.buf[1], self.buf[2], bytes(self.buf[5:5 + plen])))
            del self.buf[:total]
        return frames


def parse_rsp(payload):
    """OTA_RSP payload: {cmd_echo u8, status u8, offset u32 LE}."""
    if len(payload) < 6:
        raise RuntimeError("short OTA_RSP")
    return payload[0], payload[1], int.from_bytes(payload[2:6], "little")


# ---------------------------------------------------------------------------
# Transports: xfer(frame_bytes) -> rsp frame bytes (lockstep request/ack)
# ---------------------------------------------------------------------------

class BleTransport:
    """BLE OTA characteristic ...e0005: write frame, then read the RSP.

    The ATT write confirmation can complete before the device's BLE task has
    finished processing the write indication (and a DATA frame may trigger a
    flash sector erase, tens of ms), so the RSP read polls until non-empty."""

    def __init__(self, client):
        self.client = client

    async def xfer(self, frame):
        await self.client.write_gatt_char(OTA_RW_UUID, frame, response=True)
        deadline = asyncio.get_event_loop().time() + RSP_TIMEOUT_S
        while True:
            rsp = bytes(await self.client.read_gatt_char(OTA_RW_UUID))
            if rsp:
                return rsp
            if asyncio.get_event_loop().time() > deadline:
                raise asyncio.TimeoutError("no OTA_RSP on BLE read")
            await asyncio.sleep(0.02)


class UartTransport:
    """UART binary mode: frames on the wire, RSP frames come back on RX.

    Mode switching handshake (firmware V1.00.21+):
      text CLI --"ota"--> binary frame mode; the device prints the marker
      "[ota] binary mode ON" as the go-ahead. Exits: OTA_ABORT frame
      (immediate), 3 s idle timeout (fallback), reset after OTA_END.
    Re-sync: a bare "\\r" gets an immediate "caiic->" prompt in CLI mode;
    in binary mode it is silently dropped and the 3 s idle timeout brings
    the CLI back on its own."""

    MARKER_ON = b"[ota] binary mode ON"
    PROMPT = b"caiic->"

    def __init__(self, port, baud=115200):
        import serial
        self.ser = serial.Serial(port, baud, timeout=0.1)
        self.dec = FrameDecoder()

    def _read_until(self, markers, timeout_s):
        """Read text until any marker appears or the timeout expires.
        Runs in a worker thread - no asyncio loop time here."""
        buf = b""
        deadline = time.monotonic() + timeout_s
        while time.monotonic() < deadline:
            buf += self.ser.read(self.ser.in_waiting or 1)
            if any(m in buf for m in markers):
                break
        return buf

    async def _resync_cli(self):
        """Make sure the device sits at the text CLI (probe with CR; if a
        previous run left it in binary mode, wait out the 3 s idle exit)."""
        self.ser.reset_input_buffer()
        self.ser.write(b"\r")
        text = await asyncio.to_thread(self._read_until, [self.PROMPT], 0.8)
        if self.PROMPT in text:
            return True
        # maybe stuck in binary mode: idle timeout exits within ~3.5 s
        text = await asyncio.to_thread(self._read_until, [self.PROMPT], 3.5)
        return self.PROMPT in text

    async def enter_ota_mode(self):
        """Text CLI first: devinfo tells us the running bank ("cur bank:
        APPn"); then "ota" switches the frontend to binary frame mode,
        confirmed by the marker line (no blind sleeps)."""
        if not await self._resync_cli():
            print("warning: no CLI prompt on UART (check wiring/baud)")
        self.ser.write(b"devinfo\r")
        # read the FULL output (until the prompt) - breaking at the "cur
        # bank:" prefix can cut the line before APP1/APP2 arrives
        text = await asyncio.to_thread(self._read_until, [self.PROMPT], 2.0)
        m = re.search(rb"cur bank:\s*APP([12])", text)
        if not m:
            raise RuntimeError("could not read cur bank via CLI "
                               "(got: %r)" % text[-120:])
        bank = int(m.group(1))
        self.ser.write(b"ota\r")
        text = await asyncio.to_thread(self._read_until, [self.MARKER_ON], 2.0)
        if self.MARKER_ON not in text:
            raise RuntimeError("ota mode handshake failed (marker not seen; "
                               "firmware >= V1.00.21 required; got: %r)"
                               % text[-120:])
        self.ser.reset_input_buffer()   # drop trailing notice text
        return bank

    async def leave_ota_mode(self):
        """Explicit exit: OTA_ABORT makes the device drop back to the CLI
        immediately (also the failure-path cleanup)."""
        try:
            await self.xfer(encode_frame(TYPE_OTA_ABORT, b""))
        except Exception:
            pass
        await asyncio.to_thread(self._read_until, [self.PROMPT], 1.0)

    async def xfer(self, frame):
        self.ser.write(frame)
        raw = bytearray()
        deadline = asyncio.get_event_loop().time() + RSP_TIMEOUT_S
        while True:
            data = await asyncio.to_thread(self.ser.read, 256)
            raw += data
            for ftype, seq, payload in self.dec.feed(data):
                if ftype == TYPE_OTA_RSP:
                    # hand back a raw re-encoded frame (uniform with BLE)
                    body = bytes([ftype, seq, len(payload) & 0xFF,
                                  (len(payload) >> 8) & 0xFF]) + payload
                    crc = crc16_ccitt(body)
                    return bytes([0xCA]) + body + bytes([crc & 0xFF, crc >> 8])
            if asyncio.get_event_loop().time() > deadline:
                raise asyncio.TimeoutError("no OTA_RSP on UART (raw rx: %s)"
                                           % bytes(raw).hex(" "))


# ---------------------------------------------------------------------------
# Common helpers
# ---------------------------------------------------------------------------

def parse_combo(path):
    """Parse the combined OTA package; returns (version, {bank: payload}).

    Header (52B LE): magic | hdr_len | version | total_size | payload_crc |
    count(=2) | b1_off/size/crc | b2_off/size/crc | hdr_crc."""
    data = open(path, "rb").read()
    if len(data) < COMBO_HDR_LEN:
        raise SystemExit("bad package: too small")
    (magic, hdr_len, version, total_size, payload_crc, count,
     b1o, b1s, b1c, b2o, b2s, b2c, hcrc) = \
        struct.unpack_from("<IIIIIIIIIIIII", data, 0)
    if magic != COMBO_MAGIC or count != 2 or hdr_len != COMBO_HDR_LEN:
        raise SystemExit("bad package: magic/hdr_len/count mismatch")
    if crc32(data[:48]) != hcrc:
        raise SystemExit("bad package: header crc mismatch")
    if len(data) != total_size:
        raise SystemExit("bad package: file size %d != header total_size %d"
                         % (len(data), total_size))
    if crc32(data[hdr_len:]) != payload_crc:
        raise SystemExit("bad package: payload crc mismatch")
    banks = {}
    for bank, off, size, crc in ((1, b1o, b1s, b1c), (2, b2o, b2s, b2c)):
        blob = data[off:off + size]
        if len(blob) != size or crc32(blob) != crc:
            raise SystemExit("bad package: bank%d payload crc mismatch" % bank)
        banks[bank] = blob
    return version, banks


def fmt_ver(v):
    return "V%d.%02d.%02d" % (v >> 16, (v >> 8) & 0xFF, v & 0xFF)


BANK_BASE = {1: 0x01008000, 2: 0x01024000}


def check_blob_bank(blob, target):
    """Pre-flight: the payload's vector-table reset PC must land inside the
    target bank (the device re-checks this at OTA_END, but catching a bank
    mis-selection here saves a whole transfer)."""
    pc = struct.unpack_from("<I", blob, 4)[0]
    base = BANK_BASE[target]
    if not (base + 8 <= pc < base + len(blob)):
        raise SystemExit("payload reset PC 0x%08X is not inside APP%d "
                         "(0x%08X..0x%08X) - bank detection must be wrong, "
                         "refusing" % (pc, target, base, base + len(blob)))


def decode_tlv(data):
    items = {}
    i = 0
    while i + 2 <= len(data):
        iid, ilen = data[i], data[i + 1]
        if i + 2 + ilen > len(data):
            break
        items[iid] = data[i + 2:i + 2 + ilen]
        i += 2 + ilen
    return items


async def ota_session(xfer, chunk, blob, version, target):
    """Lockstep OTA: BEGIN -> DATA* -> END, every frame acked by OTA_RSP."""

    tty = sys.stdout.isatty()
    state = {"resent": 0, "resynced": 0}

    async def call(ftype, payload, ack_off=None):
        stage = {TYPE_OTA_BEGIN: "BEGIN", TYPE_OTA_DATA: "DATA",
                 TYPE_OTA_END: "END", TYPE_OTA_ABORT: "ABORT"}.get(ftype, "?")
        # A lost frame (flaky wiring) is recovered by resending: the device
        # re-acks a duplicate BEGIN as BAD_STATE (session already open) and a
        # duplicate DATA as BAD_STATE carrying the offset it already expects.
        for attempt in range(1, XFER_TRIES + 1):
            try:
                rsp = await xfer(encode_frame(ftype, payload))
            except asyncio.TimeoutError:
                if attempt < XFER_TRIES:
                    state["resent"] += 1
                    print("\n  %s: ack timeout, resending (attempt %d/%d)"
                          % (stage, attempt + 1, XFER_TRIES))
                    continue
                raise asyncio.TimeoutError("%s: no OTA_RSP after %d tries"
                                           % (stage, XFER_TRIES))
            frames = FrameDecoder().feed(rsp)
            if len(frames) != 1 or frames[0][0] != TYPE_OTA_RSP:
                raise RuntimeError("%s: bad RSP frame: %s" % (stage, rsp.hex(" ")))
            echo, status, offset = parse_rsp(frames[0][2])
            if echo != ftype:
                raise RuntimeError("%s: RSP echo 0x%02X != 0x%02X" % (stage, echo, ftype))
            if ftype == TYPE_OTA_BEGIN and status == ST_BAD_STATE:
                state["resynced"] += 1
                print("  BEGIN: session already open on device, resynced")
                return ST_OK, offset
            if ftype == TYPE_OTA_DATA and status == ST_BAD_STATE \
                    and ack_off is not None and offset == ack_off:
                state["resynced"] += 1
                print("\n  DATA: chunk already stored (ack was lost), resynced")
                return ST_OK, offset
            return status, offset

    def progress(sent, total, t0, rtt_ms):
        """Fine-grained progress: live-updating bar on a tty (throttled to
        ~10 fps), one line per 5%% step when output is redirected."""
        now = asyncio.get_event_loop().time()
        dt = now - t0
        pct = 100.0 * sent / total
        rate = sent / 1024.0 / max(dt, 0.001)
        if tty:
            if sent != total and now - progress.last_draw < 0.1:
                return
            progress.last_draw = now
            filled = int(pct / 100 * 24)
            bar = "#" * filled + "-" * (24 - filled)
            print("  [%s] %5.1f%%  %d/%d B  %.1f kB/s  rtt %.0f ms   "
                  % (bar, pct, sent, total, rate, rtt_ms),
                  end="\r" if sent != total else "\n", flush=True)
        else:
            step = int(pct // 5)
            if sent == total or step != progress.last_step:
                progress.last_step = step
                print("  %d / %d bytes (%.0f%%, %.1f kB/s, rtt %.0f ms)"
                      % (sent, total, pct, rate, rtt_ms))

    progress.last_draw = 0.0
    progress.last_step = -1

    print("OTA_BEGIN: open session (%d bytes, crc32=0x%08X) ..."
          % (len(blob), crc32(blob)))
    t_begin = asyncio.get_event_loop().time()
    status, _ = await call(TYPE_OTA_BEGIN,
                           struct.pack("<III", len(blob), crc32(blob), version))
    if status != 0:
        raise RuntimeError("OTA_BEGIN rejected, status=%d" % status)
    print("OTA_BEGIN ok (%.0f ms), streaming data ..."
          % ((asyncio.get_event_loop().time() - t_begin) * 1000))

    sent = 0
    t0 = asyncio.get_event_loop().time()
    while sent < len(blob):
        n = min(chunk, len(blob) - sent)
        tf = asyncio.get_event_loop().time()
        status, next_off = await call(TYPE_OTA_DATA,
                                      struct.pack("<I", sent) + blob[sent:sent + n],
                                      ack_off=sent + n)
        rtt_ms = (asyncio.get_event_loop().time() - tf) * 1000
        if status != 0:
            raise RuntimeError("OTA_DATA rejected at %d, status=%d (device "
                               "expects offset %d)" % (sent, status, next_off))
        sent += n
        progress(sent, len(blob), t0, rtt_ms)

    print("OTA_END: verify image + switch bank ...")
    try:
        status, _ = await call(TYPE_OTA_END, struct.pack("<I", crc32(blob)))
        if status != 0:
            raise RuntimeError("OTA_END rejected, status=%d "
                               "(4=crc_fail, 6=bank_mismatch)" % status)
    except (asyncio.TimeoutError, OSError) as exc:
        # the device may reset before the END write/read ack completes
        # (observed on WinRT: "operation cancelled"); the reboot wait below
        # is the real verdict
        print("OTA_END ack lost (%s) - waiting for reboot anyway ..." % exc)
    dt = asyncio.get_event_loop().time() - t0
    print("OTA_END done (%.1fs total, %d resent, %d resynced), "
          "device reboots into APP%d ..."
          % (dt, state["resent"], state["resynced"], target))
    return True


# ---------------------------------------------------------------------------
# Channel frontends
# ---------------------------------------------------------------------------

async def find_device():
    from bleak import BleakScanner
    print("scanning for %s* ..." % NAME_PREFIX)
    while True:
        dev = await BleakScanner.find_device_by_filter(
            lambda d, ad: d.name and d.name.startswith(NAME_PREFIX),
            timeout=10.0)
        if dev is not None:
            return dev
        print("  not found, retrying (device advertising?)")


async def run_ble(pkg_path):
    from bleak import BleakClient

    version, banks = parse_combo(pkg_path)
    print("package: version %s, bank1 %dB / bank2 %dB"
          % (fmt_ver(version), len(banks[1]), len(banks[2])))

    dev = await find_device()
    print("connecting %s (%s) ..." % (dev.name, dev.address))
    disconnected = asyncio.Event()

    async with BleakClient(dev, disconnected_callback=lambda _c: disconnected.set()) as client:
        print("connected, mtu=%d" % client.mtu_size)
        await asyncio.sleep(0.5)   # let the device's MTU exchange finish

        items = decode_tlv(bytes(await client.read_gatt_char(INFO_RD_UUID)))
        cur_bank = items.get(0x07, b"\x01")[0]
        cur_ver = int.from_bytes(items.get(0x01, b"\x00" * 4), "little")
        print("device: running APP%d, firmware %s" % (cur_bank, fmt_ver(cur_ver)))

        target = 3 - cur_bank
        blob = banks[target]
        check_blob_bank(blob, target)
        print("target bank: APP%d (%d bytes, crc32=0x%08X)"
              % (target, len(blob), crc32(blob)))
        if cur_ver == version:
            print("note: same version already on the device, updating anyway")

        transport = BleTransport(client)
        disconnected.clear()
        # one 0xCA frame per ATT write; data = mtu-3 minus frame overhead(7+4)
        chunk = min(client.mtu_size - 3 - 11, 222)
        try:
            await ota_session(transport.xfer, chunk, blob, version, target)
        except Exception as exc:
            print("FAIL: %s" % exc)
            try:
                await transport.xfer(encode_frame(TYPE_OTA_ABORT, b""))
            except Exception:
                pass
            return 1

        try:
            await asyncio.wait_for(disconnected.wait(), REBOOT_WAIT_S)
        except asyncio.TimeoutError:
            print("FAIL: device did not reboot within %.0fs" % REBOOT_WAIT_S)
            return 1

    print("reconnecting to verify ...")
    dev = await find_device()
    async with BleakClient(dev) as client:
        await asyncio.sleep(0.5)
        items = decode_tlv(bytes(await client.read_gatt_char(INFO_RD_UUID)))
        new_bank = items.get(0x07, b"\x00")[0]
        new_ver = int.from_bytes(items.get(0x01, b"\x00" * 4), "little")

    if new_bank == target and new_ver == version:
        print("PASS: now running APP%d, firmware %s" % (new_bank, fmt_ver(new_ver)))
        return 0
    print("FAIL: after reboot running APP%d, firmware %s (expected APP%d, %s)"
          % (new_bank, fmt_ver(new_ver), target, fmt_ver(version)))
    return 1


async def run_uart(port, pkg_path):
    version, banks = parse_combo(pkg_path)
    print("package: version %s, bank1 %dB / bank2 %dB"
          % (fmt_ver(version), len(banks[1]), len(banks[2])))

    transport = UartTransport(port)
    try:
        cur_bank = await transport.enter_ota_mode()
    except Exception as exc:
        print("FAIL: %s" % exc)
        return 1
    print("device: running APP%d" % cur_bank)
    target = 3 - cur_bank
    blob = banks[target]
    check_blob_bank(blob, target)
    print("target bank: APP%d (%d bytes, crc32=0x%08X)"
          % (target, len(blob), crc32(blob)))

    try:
        await ota_session(transport.xfer, 222, blob, version, target)
    except Exception as exc:
        print("FAIL: %s" % exc)
        await transport.leave_ota_mode()   # ABORT: device returns to the CLI
        return 1

    # The device resets after OTA_END; verify over the text CLI like the BLE
    # path verifies over GATT - the reboot result is the real verdict.
    print("waiting for reboot ...")
    await asyncio.sleep(3.0)
    transport.ser.reset_input_buffer()
    transport.ser.write(b"\r")
    await asyncio.to_thread(transport._read_until, [UartTransport.PROMPT], 2.0)
    transport.ser.write(b"devinfo\r")
    text = await asyncio.to_thread(transport._read_until,
                                   [UartTransport.PROMPT], 2.0)
    mb = re.search(rb"cur bank:\s*APP([12])", text)
    mv = re.search(rb"firmware:\s*\S+\s*\(0x([0-9A-Fa-f]{8})\)", text)
    new_bank = int(mb.group(1)) if mb else 0
    new_ver = int(mv.group(1), 16) if mv else -1
    if new_bank == target and new_ver == version:
        print("PASS: now running APP%d, firmware %s" % (new_bank, fmt_ver(new_ver)))
        return 0
    print("FAIL: after reboot running APP%d, firmware %s (expected APP%d, %s)"
          % (new_bank, fmt_ver(new_ver) if new_ver >= 0 else "?",
             target, fmt_ver(version)))
    return 1


def main():
    args = [a for a in sys.argv[1:] if not a.startswith("--")]
    uart = None
    if "--uart" in sys.argv[1:]:
        i = sys.argv[1:].index("--uart")
        uart = sys.argv[1:][i + 1]
        args = [a for j, a in enumerate(sys.argv[1:]) if j not in (i, i + 1)]

    root = os.path.dirname(os.path.abspath(__file__))
    pkg = args[0] if args else os.path.join(root, "out", "mothercup_ble_ota.bin")
    if not os.path.isfile(pkg):
        print("package not found: %s (run tools\\make_package.bat first)" % pkg)
        return 1
    if uart:
        return asyncio.run(run_uart(uart, pkg))
    return asyncio.run(run_ble(pkg))


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        pass