#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
ble_cli_test.py - PC-side BLE test client for the CAIIC-MCM frame protocol
(V1.1), used to verify the mcm-ddc-ble firmware downlink/uplink without a
phone in the loop.

    venv:   .venv-ble  (python -m venv .venv-ble && pip install bleak)
    usage:  .venv-ble/Scripts/python ble_cli_test.py            # framed CLI: send "help"
            .venv-ble/Scripts/python ble_cli_test.py "sysinfo"  # framed CLI: one command
            .venv-ble/Scripts/python ble_cli_test.py -i         # interactive framed mode
            .venv-ble/Scripts/python ble_cli_test.py -rw help   # CLI via the new read/write
                                                                #   characteristic ...0003
            .venv-ble/Scripts/python ble_cli_test.py -rw -i     # interactive rw mode
            in framed interactive mode: "info" sends an INFO_QUERY (all items)

Watch the device UART (blelog on) in parallel: a healthy exchange prints
[BLE] notify enabled / [BLE] write ind: att_idx=2 ... / [BLE RX] / [BLE TX].
"""

import asyncio
import sys

from bleak import BleakClient, BleakScanner

NAME_PREFIX = "CAIIC-MCM-20260902"
SERVICE_UUID = "00002760-08c2-11e1-9073-0e8ac72e1001"
WRITE_UUID = "00002760-08c2-11e1-9073-0e8ac72e0001"   # Write Without Response (framed protocol)
NOTIFY_UUID = "00002760-08c2-11e1-9073-0e8ac72e0002"  # Notify
CLI_RW_UUID = "00002760-08c2-11e1-9073-0e8ac72e0003"  # Read + Write: raw CLI line in, output text out
INFO_RD_UUID = "00002760-08c2-11e1-9073-0e8ac72e0004"  # Read-only: device info TLV (all items)

TYPE_CLI_REQ = 0x01
TYPE_CLI_RSP = 0x02
TYPE_CLI_RSP_END = 0x03
TYPE_INFO_QUERY = 0x20
TYPE_INFO_RSP = 0x21

FRAME_SOF = 0xCA
MAX_PAYLOAD = 480

CLI_TIMEOUT_S = 5.0
INFO_TIMEOUT_S = 3.0


def crc16_ccitt(data):
    """CRC16-CCITT, poly 0x1021, init 0xFFFF (same as firmware proto_crc16)."""
    crc = 0xFFFF
    for b in data:
        crc ^= b << 8
        for _ in range(8):
            if crc & 0x8000:
                crc = ((crc << 1) ^ 0x1021) & 0xFFFF
            else:
                crc = (crc << 1) & 0xFFFF
    return crc


def encode_frame(ftype, seq, payload):
    """SOF + TYPE + SEQ + LEN(u16 LE) + PAYLOAD + CRC16(over TYPE..PAYLOAD, LE)."""
    out = bytearray([FRAME_SOF, ftype & 0xFF, seq & 0xFF,
                     len(payload) & 0xFF, (len(payload) >> 8) & 0xFF])
    out += bytes(payload)
    crc = crc16_ccitt(out[1:])
    out += bytes([crc & 0xFF, (crc >> 8) & 0xFF])
    return bytes(out)


class FrameDecoder:
    """Byte-stream reassembly, mirrors firmware app_ble_proto_rx_bytes."""

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

    def feed(self, data):
        frames = []
        self.buf += bytes(data)
        while True:
            while self.buf and self.buf[0] != FRAME_SOF:
                del self.buf[0]
            if len(self.buf) < 5:
                break
            plen = self.buf[3] | (self.buf[4] << 8)
            if plen > MAX_PAYLOAD:
                del self.buf[0]
                continue
            total = 5 + plen + 2
            if len(self.buf) < total:
                break
            crc = crc16_ccitt(self.buf[1:5 + plen])
            rx_crc = self.buf[5 + plen] | (self.buf[5 + plen + 1] << 8)
            if crc != rx_crc:
                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


class McmTestClient:
    def __init__(self):
        self.decoder = FrameDecoder()
        self.tx_seq = 0
        self.cli_chunks = []
        self.cli_done = None        # asyncio.Event, set on CLI_RSP_END
        self.cli_status = -1
        self.info_done = None       # asyncio.Event, set on INFO_RSP
        self.info_payload = None

    def on_notify(self, _sender, data):
        frames = self.decoder.feed(data)
        for ftype, seq, payload in frames:
            print("[frame] type=0x%02X seq=%u payload=%s"
                  % (ftype, seq, payload.hex(" ")))
            if ftype == TYPE_CLI_RSP:
                self.cli_chunks.append(payload.decode("ascii", "replace"))
            elif ftype == TYPE_CLI_RSP_END:
                self.cli_status = payload[0] if payload else 1
                if self.cli_done:
                    self.cli_done.set()
            elif ftype == TYPE_INFO_RSP:
                self.info_payload = payload
                if self.info_done:
                    self.info_done.set()

    async def send_frame(self, client, ftype, payload):
        frame = encode_frame(ftype, self.tx_seq, payload)
        self.tx_seq = (self.tx_seq + 1) & 0xFF
        print("[tx] %s" % frame.hex(" "))
        # downlink characteristic is Write Without Response only
        await client.write_gatt_char(WRITE_UUID, frame, response=False)

    async def run_cli(self, client, cmd):
        payload = cmd.encode("ascii")
        if len(payload) > 63:
            print("command too long (>63 bytes)")
            return
        self.cli_chunks = []
        self.cli_status = -1
        self.cli_done = asyncio.Event()
        await self.send_frame(client, TYPE_CLI_REQ, payload)
        try:
            await asyncio.wait_for(self.cli_done.wait(), CLI_TIMEOUT_S)
        except asyncio.TimeoutError:
            print("[rx] TIMEOUT - no CLI_RSP_END within %.0fs" % CLI_TIMEOUT_S)
            return
        text = "".join(self.cli_chunks)
        print("[rx] status=%d%s" % (self.cli_status, ("\n" + text) if text else " (ok)"))

    async def run_info_query(self, client):
        self.info_payload = None
        self.info_done = asyncio.Event()
        await self.send_frame(client, TYPE_INFO_QUERY, b"")
        try:
            await asyncio.wait_for(self.info_done.wait(), INFO_TIMEOUT_S)
        except asyncio.TimeoutError:
            print("[rx] TIMEOUT - no INFO_RSP within %.0fs" % INFO_TIMEOUT_S)
            return
        print("[rx] INFO_RSP TLV: %s" % self.info_payload.hex(" "))


async def run_cli_rw(client, cmd):
    """CLI over the read/write characteristic ...0003: write the raw command
    line (ATT write-with-response), then read the output text back in
    sequential chunks (a 0-length read marks the end of the response)."""
    print("[tx rw] %r" % cmd)
    await client.write_gatt_char(CLI_RW_UUID, cmd.encode("ascii"), response=True)
    await asyncio.sleep(0.1)  # one conn. interval for the device to execute
    rsp = bytearray()
    for _ in range(64):  # safety bound; each chunk <= att_mtu-2
        chunk = await client.read_gatt_char(CLI_RW_UUID)
        if not chunk:
            break
        rsp += chunk
    print("[rx rw] %dB:\n%s" % (len(rsp), rsp.decode("ascii", "replace")))


_INFO_ITEM_NAMES = {0x01: "fw_version", 0x02: "chip_temp", 0x03: "fan_rpm",
                    0x04: "vdd_mv", 0x05: "uptime_s", 0x06: "free_heap",
                    0x07: "cur_bank"}


async def run_param_read(client):
    """Read the read-only device-info characteristic ...0004 and decode the
    TLV sequence {id u8, len u8, value LE}*n."""
    data = bytes(await client.read_gatt_char(INFO_RD_UUID))
    print("[rx rd] %dB TLV: %s" % (len(data), data.hex(" ")))
    i = 0
    while i + 2 <= len(data):
        iid, ilen = data[i], data[i + 1]
        if i + 2 + ilen > len(data):
            break
        raw = data[i + 2:i + 2 + ilen]
        name = _INFO_ITEM_NAMES.get(iid, "0x%02X" % iid)
        val = int.from_bytes(raw, "little")
        if iid == 0x01 and ilen == 4:
            val = "V%d.%02d.%02d" % (val >> 16, (val >> 8) & 0xFF, val & 0xFF)
        elif iid == 0x02 and ilen == 2:
            t = int.from_bytes(raw, "little", signed=True)
            val = "invalid" if t == 0x7FFF else "%.1f C" % (t / 10.0)
        elif iid == 0x03 and ilen == 2 and val == 0xFFFF:
            val = "n/a"
        print("  %-10s %s" % (name, val))
        i += 2 + ilen


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? phone disconnected?)")


async def main():
    args = sys.argv[1:]
    interactive = "-i" in args
    use_rw = "-rw" in args
    cmds = [a for a in args if a not in ("-i", "-rw")] or ["help"]

    dev = await find_device()
    print("connecting %s (%s) ..." % (dev.name, dev.address))
    async with BleakClient(dev) as client:
        print("connected, mtu=%d" % client.mtu_size)
        tester = McmTestClient()
        if not use_rw:
            await client.start_notify(NOTIFY_UUID, tester.on_notify)
            print("notify subscribed on %s" % NOTIFY_UUID)
        # give the device a moment; it also initiates its own MTU exchange
        await asyncio.sleep(0.5)
        if interactive:
            print("interactive mode (%s), type a CLI command ('exit'):"
                  % ("rw-char, 'param' = read info char" if use_rw
                     else "framed, 'info' = INFO_QUERY, 'param' = read info char"))
            while True:
                line = (await asyncio.to_thread(input, "ble> ")).strip()
                if line in ("exit", "quit"):
                    break
                if not line:
                    continue
                if use_rw:
                    if line == "param":
                        await run_param_read(client)
                    else:
                        await run_cli_rw(client, line)
                elif line == "info":
                    await tester.run_info_query(client)
                elif line == "param":
                    await run_param_read(client)
                else:
                    await tester.run_cli(client, line)
        else:
            # one-shot: the whole argv is ONE command line (multi-word ok)
            line = " ".join(cmds)
            if line == "param":
                await run_param_read(client)
            elif use_rw:
                await run_cli_rw(client, line)
            elif line == "info":
                await tester.run_info_query(client)
            else:
                await tester.run_cli(client, line)
        if not use_rw:
            await client.stop_notify(NOTIFY_UUID)
    print("disconnected")


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        pass