#!/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 (including `watch gcp` output), 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}")

        print(f"{'watch gcp':<12} -> {await command(client, replies, 'watch gcp')}")
        records = [json.loads(await asyncio.wait_for(replies.lines.get(), timeout=3)) for _ in range(2)]
        await client.write_gatt_char(COMMAND, b"watch off", response=True)
        await asyncio.sleep(1.5)
        while not replies.lines.empty():
            replies.lines.get_nowait()  # the watch reply and any second pushed before it
        print(f"{'gcp seconds':<12} -> {[r['second'] for r in records]}, healthy {records[-1]['healthy']}, "
              f"sample sums {records[-1]['sample']['sums']}")

        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())