#!/usr/bin/env python3
"""End-to-end check of the STRIX firmware's USB API (arduino/strix, see API.md).

Runs every command once, checks the replies, measures entropy throughput, and runs the onod3000
suite on 256 KB from the entropy stream.
"""

import argparse
import json
import subprocess
import sys
import time
from pathlib import Path

import serial


class Device:
    def __init__(self, port: str):
        self.serial = serial.Serial(port, 115200, timeout=5)
        time.sleep(0.5)
        self.serial.reset_input_buffer()

    def line(self, command: str) -> str:
        self.serial.write(f"{command}\n".encode())
        return self.serial.readline().decode().strip()

    def json(self, command: str) -> dict:
        reply = self.line(command)
        try:
            return json.loads(reply)
        except json.JSONDecodeError:
            sys.exit(f"{command!r}: not JSON: {reply!r}")

    def entropy_bin(self, count: int) -> bytes:
        self.serial.write(f"entropy {count} bin\n".encode())
        header = self.serial.readline().decode().strip()
        if header != f"BIN {count}":
            sys.exit(f"entropy bin: unexpected header {header!r}")
        data = self.serial.read(count)
        if len(data) != count:
            sys.exit(f"entropy bin: {len(data)} of {count} bytes")
        return data


def check(condition: bool, message: str) -> None:
    print(("ok    " if condition else "FAIL  ") + message)
    if not condition:
        check.failures += 1


check.failures = 0


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--port", default="/dev/ttyACM0")
    parser.add_argument("--onod", type=Path, default=Path(".build/onod3000/release/onod3000"))
    args = parser.parse_args()
    device = Device(args.port)

    info = device.json("id")
    check(info.get("device") == "STRIX" and info.get("fpga_ready"), f"id: {info}")

    health = device.json("health")
    check(health.get("healthy") == 15, f"health: all 4 generators healthy (mask {health.get('healthy')})")

    config = device.json("get")
    check(config.get("temporal") == 5 and config.get("stream_source") == 4, f"get: {config}")

    hex_line = device.line("entropy 32")
    check(len(hex_line) == 64 and all(c in "0123456789ABCDEF" for c in hex_line), f"entropy 32 hex: {hex_line}")

    start = time.time()
    device.entropy_bin(65536)
    elapsed = time.time() - start
    check(True, f"entropy 65536 bin: {65536 / elapsed / 1024:.0f} KB/s")

    check(device.json("set rolling 5") == {"ok": True}, "set rolling 5")
    check(device.json("get").get("rolling") == 5, "get rolling == 5")
    check("error" in device.json("set temporal 9"), "set temporal 9 is refused")
    device.json("set rolling 10")

    device.json("stats reset")
    time.sleep(1.2)
    stats = device.json("stats")
    channels = stats.get("channels", [])
    check(len(channels) == 5 and stats.get("trials", 0) >= 4 and stats.get("missed") == 0,
          f"stats: {stats.get('trials')} trials, {stats.get('missed')} missed, "
          f"combined total z {channels[-1]['total']['z'] if channels else None}")

    amp = device.json("amp")
    check(len(amp.get("paths", [])) == 2 and not amp.get("overflow"),
          f"amp: trial {amp.get('trial')}, bias {amp['paths'][0]['bits']} bits, votes {amp['paths'][0]['votes']}")

    check(device.json("watch stats") == {"ok": True}, "watch stats")
    trials = [json.loads(device.serial.readline()).get("trial") for _ in range(3)]
    device.json("watch off")
    device.serial.reset_input_buffer()
    check(trials == sorted(trials) and len(set(trials)) == 3, f"watch stats pushed trials {trials}")

    check(device.json("watch gcp") == {"ok": True}, "watch gcp")
    seconds = [json.loads(device.serial.readline()) for _ in range(3)]
    device.json("watch off")
    device.serial.reset_input_buffer()
    numbers = [r.get("second") for r in seconds]
    last = seconds[-1]
    check(numbers == list(range(numbers[0], numbers[0] + 3)), f"watch gcp pushed consecutive seconds {numbers}")
    check(last["healthy"] == 15 and last["sample"]["bits"] == 200
          and all(0 <= s <= 200 for s in last["sample"]["sums"]),
          f"gcp sample: sums {last['sample']['sums']}, Stouffer {last['sample']['stouffer']}")
    check(abs(last["full"]["bits"] - 24_000_000) <= 1,
          f"gcp full: {last['full']['bits']} bits, z {last['full']['z']}, Stouffer {last['full']['stouffer']}")
    device.entropy_bin(65536)
    time.sleep(1.1)
    records = device.json(f"gcp {numbers[-1]}").get("records", [])
    after = [r["second"] for r in records]
    check(len(after) >= 1 and after == list(range(numbers[-1] + 1, numbers[-1] + 1 + len(after))),
          f"gcp {numbers[-1]}: no gap across an entropy transfer, seconds {after}")

    if args.onod.exists():
        sample = Path("captures/api_entropy.bin")
        sample.parent.mkdir(exist_ok=True)
        sample.write_bytes(b"".join(device.entropy_bin(65536) for _ in range(4)))
        report = subprocess.run([str(args.onod), "-f", str(sample)], capture_output=True, text=True).stdout
        summary = next((l for l in report.splitlines() if "tests passed" in l), "no summary")
        passed = int(summary.split("/")[0]) if "/" in summary else 0
        # 15 tests at p < 0.01: one chance failure happens in about 14% of runs on ideal data.
        check(passed >= 14, f"onod3000 on 256 KB from the stream: {summary}")
    else:
        print(f"skip  onod3000 ({args.onod} not built)")

    print(f"\n{check.failures} failure(s)")
    sys.exit(1 if check.failures else 0)


if __name__ == "__main__":
    main()