#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
ble_ota_update.py - PC-side BLE OTA updater for CAIIC-MCM devices.

Takes the combined OTA package (mothercup_ble_ota.bin: 52-byte header + both
bank-linked payloads, see tools/merge_image.py), reads the device's current
bank from the read-only info characteristic ...e0004 (info item 0x07
CUR_BANK), and downloads the payload linked for the OPPOSITE bank over the
framed protocol downlink (characteristic ...e0001).

    usage:  ble_ota.bat [combo_pkg]
            default pkg: out/mothercup_ble_ota.bin

Flow-control note: the notify uplink is currently broken (dev log section
28), so OTA_RSP acks cannot be received. The script therefore paces writes:
a short gap per frame and a longer gap at every 4KB sector boundary (the
device erases sectors lazily with interrupts off - writes during the erase
window would be lost). Success is detected by the device resetting after
OTA_END (link drops), then re-checking CUR_BANK/version after reconnect.
The firmware independently rejects a payload linked for the wrong bank
(OTA_RSP status 6, no switch), so a wrong pick never bricks the device.
"""

import asyncio
import os
import struct
import sys
import zlib

from bleak import BleakClient, BleakScanner

NAME_PREFIX = "CAIIC-MCM-20260902"
WRITE_UUID = "00002760-08c2-11e1-9073-0e8ac72e0001"   # Write Without Response
INFO_RD_UUID = "00002760-08c2-11e1-9073-0e8ac72e0004"  # Read-only info TLV

TYPE_OTA_BEGIN = 0x10
TYPE_OTA_DATA = 0x11
TYPE_OTA_END = 0x12

COMBO_MAGIC = 0xCA10BA11
COMBO_HDR_LEN = 52

FRAME_GAP_S = 0.008      # between WWR frames
SECTOR_GAP_S = 0.25      # at 4KB boundaries (lazy sector erase, IRQs off)
REBOOT_WAIT_S = 20.0
SECTOR = 0x1000


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


_seq = [0]


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)


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


async def read_info_tlv(client):
    """Read ...e0004 and decode the TLV into {item_id: raw_bytes}."""
    data = bytes(await client.read_gatt_char(INFO_RD_UUID))
    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


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


async def find_device():
    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 ota_session(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])))

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

    disconnected = asyncio.Event()

    def on_disconnect(_c):
        disconnected.set()

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

        items = await read_info_tlv(client)
        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      # OTA always writes the inactive bank
        blob = banks[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")

        max_data = min(client.mtu_size - 3 - 4 - 7, 233)   # offset u32 + frame overhead
        disconnected.clear()

        async def send(ftype, payload):
            await client.write_gatt_char(WRITE_UUID, encode_frame(ftype, payload),
                                         response=False)
            await asyncio.sleep(FRAME_GAP_S)

        for attempt in (1, 2):
            sector_gap = SECTOR_GAP_S * attempt
            print("OTA_BEGIN ... (attempt %d)" % attempt)
            await send(TYPE_OTA_BEGIN, struct.pack("<III", len(blob), crc32(blob), version))

            sent = 0
            t0 = asyncio.get_event_loop().time()
            while sent < len(blob):
                first_of_sector = (sent % SECTOR) == 0
                n = min(max_data, len(blob) - sent, SECTOR - (sent % SECTOR))
                await send(TYPE_OTA_DATA, struct.pack("<I", sent) + blob[sent:sent + n])
                sent += n
                if sent % SECTOR == 0 or sent == len(blob):
                    print("  %d / %d bytes (%.0f%%)"
                          % (sent, len(blob), 100.0 * sent / len(blob)))
                # the first frame touching a NEW sector triggers its lazy
                # erase (IRQs off, tens of ms) - frames sent during the erase
                # are lost, so wait right after that first frame
                if first_of_sector and sent < len(blob):
                    await asyncio.sleep(sector_gap)

            print("OTA_END (crc32=0x%08X), waiting for device reboot ..."
                  % crc32(blob))
            await send(TYPE_OTA_END, struct.pack("<I", crc32(blob)))

            try:
                await asyncio.wait_for(disconnected.wait(), REBOOT_WAIT_S)
                break
            except asyncio.TimeoutError:
                if attempt == 2:
                    print("FAIL: device did not reboot - OTA rejected "
                          "(crc/bank mismatch?). Check UART with 'blelog on'.")
                    return 1
                print("  no reboot; retrying the whole image "
                      "(frames were probably lost in an erase window)")
        else:
            return 1

    dt = asyncio.get_event_loop().time() - t0
    print("device rebooted (transfer %.1fs). verifying new bank ..." % dt)

    dev = await find_device()
    async with BleakClient(dev) as client:
        await asyncio.sleep(0.5)
        items = await read_info_tlv(client)
        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


def main():
    root = os.path.dirname(os.path.abspath(__file__))
    pkg = sys.argv[1] if len(sys.argv) > 1 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
    return asyncio.run(ota_session(pkg))


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