#!/usr/bin/env python3
"""Writes STRIX entropy to stdout continuously, for piping into test suites
(e.g. `python3 tools/stream_entropy.py | qngmeter`). Uses the firmware's `entropy <n> bin`
command (arduino/strix/API.md)."""

import argparse
import sys
import time

import serial

CHUNK = 65536


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--port", default="/dev/ttyACM0")
    parser.add_argument("--bytes", type=int, default=0, help="stop after this many bytes (0 = run until killed)")
    args = parser.parse_args()

    device = serial.Serial(args.port, 115200, timeout=5)
    time.sleep(0.5)
    device.reset_input_buffer()
    out = sys.stdout.buffer
    sent = 0
    try:
        while args.bytes == 0 or sent < args.bytes:
            count = CHUNK if args.bytes == 0 else min(CHUNK, args.bytes - sent)
            device.write(f"entropy {count} bin\n".encode())
            header = device.readline().decode(errors="replace").strip()
            if header != f"BIN {count}":
                sys.exit(f"device: {header or 'timeout'}")
            data = device.read(count)
            if len(data) != count:
                sys.exit(f"short read: {len(data)} of {count} bytes")
            out.write(data)
            out.flush()
            sent += count
    except (BrokenPipeError, KeyboardInterrupt):
        pass


if __name__ == "__main__":
    main()