#!/usr/bin/env python3
"""Command-line client for the STRIX firmware over USB (arduino/strix, see API.md).

Subcommands:
  id, health, stats, amp, get   print the device's reply
  set KEY VALUE                 change a setting
  entropy N                     N entropy bytes as hex, or binary into a file with --out
  watch stats|amp               print every new trial until Ctrl-C
  gcp                           follow the per-second GCP-format trials: live cumulative
                                Z^2 - 1 against its p = 0.05 envelope, optional CSV log
  send LINE                     send any command line and print the reply

The port is found by the board's USB ID unless --port or STRIX_PORT is given.
"""

import argparse
import json
import math
import os
import sys
import time
from collections import deque
from datetime import datetime, timezone

import serial
from serial.tools import list_ports

VIDOR_USB_ID = (0x2341, 0x8056)
MAX_ENTROPY_REQUEST = 65536
ENVELOPE_Z = 1.96  # two-sided p = 0.05


def find_port() -> str:
    if os.environ.get("STRIX_PORT"):
        return os.environ["STRIX_PORT"]
    for port in list_ports.comports():
        if (port.vid, port.pid) == VIDOR_USB_ID:
            return port.device
    return "/dev/ttyACM0"


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

    def send(self, line: str) -> None:
        self.serial.write(f"{line}\n".encode())

    def read_line(self) -> str:
        line = self.serial.readline()
        if not line:
            raise TimeoutError("no reply from the device")
        return line.decode(errors="replace").strip()

    def json(self, line: str) -> dict:
        self.send(line)
        return json.loads(self.read_line())

    def entropy(self, count: int) -> bytes:
        self.send(f"entropy {count} bin")
        header = self.read_line()
        if header != f"BIN {count}":
            raise RuntimeError(header)  # an error reply such as an unhealthy stream source
        data = self.serial.read(count)
        if len(data) != count:
            raise RuntimeError(f"stream stopped after {len(data)} of {count} bytes")
        return data


def print_reply(reply: dict) -> None:
    print(json.dumps(reply, indent=2))
    if "error" in reply:
        sys.exit(1)


def run_entropy(device: Device, args) -> None:
    out = open(args.out, "wb") if args.out else None
    remaining = args.count
    while remaining:
        chunk = device.entropy(min(remaining, MAX_ENTROPY_REQUEST))
        if out:
            out.write(chunk)
        else:
            sys.stdout.write(chunk.hex().upper())
        remaining -= len(chunk)
    if out:
        out.close()
    else:
        print()


def run_watch(device: Device, args) -> None:
    reply = device.json(f"watch {args.what}")
    if "error" in reply:
        print_reply(reply)
    try:
        while True:
            print(device.read_line(), flush=True)
    except KeyboardInterrupt:
        device.send("watch off")


class GcpFollower:
    """Turns per-second records into timestamped trials and a cumulative network variance.

    The device counts seconds since its FPGA started but has no clock, so the host anchors
    second numbers to its own clock as live records arrive; records fetched to fill a gap get
    times from the same anchor."""

    def __init__(self, measure: str, log_path: str | None):
        self.measure = measure
        self.last_second = None
        self.offset = None  # host time at the end of second 0
        self.trials = 0
        self.cumulative = 0.0
        self.log = None
        if log_path:
            new = not os.path.exists(log_path) or os.path.getsize(log_path) == 0
            self.log = open(log_path, "a")
            if new:
                self.log.write("12,gmtime,Date/Time,g0,g1,g2,g3\n")

    def anchor(self, second: int) -> None:
        """Called with each live record, which arrives just after its second ended."""
        now = time.time()
        if self.offset is None or abs(now - (self.offset + second)) > 2:
            self.offset = math.floor(now) - second

    def handle(self, record: dict) -> None:
        second = record["second"]
        if self.last_second is not None and second <= self.last_second:
            print(f"device restarted (second {second} after {self.last_second})", file=sys.stderr)
        elif self.last_second is not None and second > self.last_second + 1:
            print(f"{second - self.last_second - 1} second(s) lost", file=sys.stderr)
        self.last_second = second

        start = self.offset + second - 1
        stamp = datetime.fromtimestamp(start, timezone.utc).strftime("%Y-%m-%d %H:%M:%S")
        sums = record["sample"]["sums"]
        if self.log:
            fields = ["" if s is None else str(s) for s in sums]
            self.log.write(f"13,{start},{stamp},{','.join(fields)}\n")
            self.log.flush()

        trial = record[self.measure]
        netvar = trial["netvar"]
        if netvar is None:
            print(f"{stamp}  #{second:<7} no healthy generator")
            return
        self.trials += 1
        self.cumulative += netvar - 1
        envelope = ENVELOPE_Z * math.sqrt(2 * self.trials)
        outside = " *" if abs(self.cumulative) > envelope else ""
        shown = " ".join("  -" if s is None else f"{s:3d}" for s in sums)
        print(f"{stamp}  #{second:<7} sums {shown}  Z {trial['stouffer']:+.3f}  Z² {netvar:6.3f}"
              f"  Σ(Z²−1) {self.cumulative:+8.2f}  ±{envelope:6.2f}{outside}", flush=True)


def fetch_missed(device: Device, after: int, before: int, pending: deque) -> list[dict]:
    """Seconds after `after` and before `before` from the device's 60-second history. Live
    records arriving meanwhile are queued in `pending`."""
    device.send(f"gcp {after}")
    while True:
        line = device.read_line()
        if line.startswith('{"records"'):
            return [r for r in json.loads(line)["records"] if r["second"] < before]
        if line.startswith('{"second"'):
            pending.append(line)


def follow_gcp(port: str, follower: GcpFollower, seconds: int | None) -> None:
    """Follows `watch gcp`, reconnecting after USB dropouts and filling gaps from the device's
    history."""
    end = time.time() + seconds if seconds else None
    device = None
    pending = deque()
    try:
        while end is None or time.time() < end:
            try:
                if device is None:
                    device = Device(port)
                    reply = device.json("watch gcp")
                    if "error" in reply:
                        print_reply(reply)
                line = pending.popleft() if pending else device.read_line()
                if not line.startswith('{"second"'):
                    continue
                record = json.loads(line)
                second, last = record["second"], follower.last_second
                follower.anchor(second)
                if last is not None and second > last + 1:
                    for missed in fetch_missed(device, last, second, pending):
                        follower.handle(missed)
                follower.handle(record)
            except (serial.SerialException, OSError, TimeoutError) as error:
                print(f"connection lost ({error}), reconnecting", file=sys.stderr)
                if device:
                    device.serial.close()
                device = None
                pending.clear()
                time.sleep(2)
    except KeyboardInterrupt:
        pass
    finally:
        if device:
            device.send("watch off")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--port", help="serial port (default: the STRIX board, or STRIX_PORT)")
    commands = parser.add_subparsers(dest="command", required=True)
    for name in ("id", "health", "amp", "get"):
        commands.add_parser(name)
    stats = commands.add_parser("stats")
    stats.add_argument("--reset", action="store_true", help="restart trials and windows")
    setting = commands.add_parser("set")
    setting.add_argument("key")
    setting.add_argument("value")
    entropy = commands.add_parser("entropy")
    entropy.add_argument("count", type=int)
    entropy.add_argument("--out", help="write binary to this file instead of hex to stdout")
    watch = commands.add_parser("watch")
    watch.add_argument("what", choices=["stats", "amp"])
    gcp = commands.add_parser("gcp")
    gcp.add_argument("--log", help="append timestamped 200-bit trial sums to this CSV file")
    gcp.add_argument("--full", action="store_true", help="track the full-rate trials instead of the 200-bit sample")
    gcp.add_argument("--seconds", type=int, help="stop after this many seconds")
    send = commands.add_parser("send")
    send.add_argument("line", nargs="+")
    args = parser.parse_args()

    port = args.port or find_port()
    if args.command == "gcp":
        follow_gcp(port, GcpFollower("full" if args.full else "sample", args.log), args.seconds)
        return

    device = Device(port)
    if args.command == "stats":
        print_reply(device.json("stats reset" if args.reset else "stats"))
    elif args.command == "set":
        print_reply(device.json(f"set {args.key} {args.value}"))
    elif args.command == "entropy":
        run_entropy(device, args)
    elif args.command == "watch":
        run_watch(device, args)
    elif args.command == "send":
        device.send(" ".join(args.line))
        print(device.read_line())
    else:
        print_reply(device.json(args.command))


if __name__ == "__main__":
    main()