#!/usr/bin/env python3
"""Checks the STRIX firmware's Bluetooth LE API (arduino/strix/API.md) from this computer:
finds the device, sends commands over the command characteristic, reads the replies from the
reply notifications, and reads and streams the entropy characteristic."""

import argparse
import asyncio
import json
import sys

from bleak import BleakClient, BleakScanner

SERVICE = "53545249-5800-4000-8000-000000000001"
COMMAND = "53545249-5800-4000-8000-000000000002"
REPLY = "53545249-5800-4000-8000-000000000003"
ENTROPY = "53545249-5800-4000-8000-000000000004"


class Replies:
    """Reassembles reply notifications into newline-terminated lines."""

    def __init__(self):
        self.buffer = b""
        self.lines: asyncio.Queue[str] = asyncio.Queue()

    def handle(self, _sender, data: bytearray) -> None:
        self.buffer += bytes(data)
        while b"\n" in self.buffer:
            line, self.buffer = self.buffer.split(b"\n", 1)
            self.lines.put_nowait(line.decode(errors="replace").strip())


async def command(client: BleakClient, replies: Replies, text: str) -> str:
    await client.write_gatt_char(COMMAND, text.encode(), response=True)
    return await asyncio.wait_for(replies.lines.get(), timeout=10)


async def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--timeout", type=float, default=15, help="scan time in seconds")
    args = parser.parse_args()

    device = await BleakScanner.find_device_by_filter(
        lambda d, adv: (d.name or "") == "STRIX" or SERVICE in adv.service_uuids, timeout=args.timeout)
    if device is None:
        sys.exit("no STRIX device found")
    print(f"found {device.name} at {device.address}")

    async with BleakClient(device) as client:
        replies = Replies()
        await client.start_notify(REPLY, replies.handle)

        for text in ["id", "health", "ble", "entropy 32"]:
            reply = await command(client, replies, text)
            print(f"{text:<12} -> {reply[:160]}")
        stats = json.loads(await command(client, replies, "stats"))
        print(f"{'stats':<12} -> trial {stats.get('trial')}, combined total z "
              f"{stats['channels'][-1]['total']['z'] if 'channels' in stats else None}")

        first = await client.read_gatt_char(ENTROPY)
        second = await client.read_gatt_char(ENTROPY)
        print(f"entropy read -> {len(first)} bytes, then {len(second)} bytes, "
              f"{'different' if first != second else 'IDENTICAL'}")

        chunks = []
        await client.start_notify(ENTROPY, lambda _s, data: chunks.append(bytes(data)))
        await asyncio.sleep(3)
        await client.stop_notify(ENTROPY)
        total = sum(len(c) for c in chunks)
        print(f"entropy stream -> {len(chunks)} notifications, {total} bytes in 3 s "
              f"({total / 3:.0f} B/s), {len(set(chunks))} distinct")


if __name__ == "__main__":
    asyncio.run(main())