#!/usr/bin/env python3
"""
max_command_speed_test.py

Finds the maximum stable command rate for the UART command protocol described
by `struct command_message_t`, using the LED command (sent with length=0,
i.e. no data bytes -- only the 5-byte header) as the test payload.

Wire format (matches the firmware's ring-buffer parser):

    [prefix:1][length:1][id:1][command:1][crc:1][data:length]

Note the firmware only ever reads `length` bytes of data off the wire (see
`usb_thread`), NOT the full 160-byte `data[]` array from the struct -- the
160 bytes only exist in RAM. So we must only transmit `length` bytes.

CRC: sum of prefix+length+id+command+0(crc placeholder)+data[0..length-1],
mod 256, then crc = (0x100 - sum) & 0xff. This matches
`command_calculate_crc()`, which runs while msg->crc is still 0 (it's set
by the caller only *after* this function returns).

Two test modes:

  roundtrip  - send one command, wait for ACK/NACK, repeat. Measures the
               sustainable rate when the host waits for each reply
               (safest / simplest way to drive the device).

  flood      - send a whole burst of commands back-to-back at a fixed
               inter-command delay, then collect all the ACK/NACK replies
               afterwards. This stresses the UART IRQ handler + ring buffer
               directly and finds the max raw throughput the firmware can
               absorb without dropping/desyncing bytes.

The script sweeps delays (roundtrip) or a binary search on delay (flood) to
find the smallest stable delay (i.e. highest command rate) that still gets
a correct ACK for every command sent.

Requires: pip install pyserial
"""

import argparse
import struct
import sys
import time

import serial

# ---- Protocol constants (from firmware) -----------------------------------
COMMAND_PREFIX = 0x69
COMMAND_ID = 0x00
COMMAND_DATA_SIZE = 160

COMMAND_ACK = 0
COMMAND_NACK = 1
LED = 2

HEADER_FMT = "<BBBBB"  # prefix, length, id, command, crc
HEADER_LEN = struct.calcsize(HEADER_FMT)


def calc_crc(length: int, command: int, data: bytes) -> int:
    """Replicates command_calculate_crc(): sum over
    [prefix, length, id, command, crc(=0 at calc time), data[:length]]"""
    s = COMMAND_PREFIX + length + COMMAND_ID + command + 0
    s += sum(data[:length])
    return (0x100 - (s & 0xFF)) & 0xFF


def build_packet(command: int, data: bytes = b"") -> bytes:
    length = len(data)
    if length > COMMAND_DATA_SIZE - 1:
        raise ValueError(f"data too long: {length} > {COMMAND_DATA_SIZE - 1}")
    crc = calc_crc(length, command, data)
    return struct.pack(HEADER_FMT, COMMAND_PREFIX, length, COMMAND_ID, command, crc) + data


def build_led_packet() -> bytes:
    """LED command with no payload (length=0). Only the 5-byte header
    [prefix][length=0][id][command][crc] is sent -- no data bytes at all."""
    return build_packet(LED, b"")


class ResponseError(Exception):
    pass


def read_response(ser: serial.Serial, timeout: float):
    """Read one full response packet (ACK/NACK) from the wire, honoring
    `timeout` seconds total. Returns dict with command/length/data/crc_ok,
    or None on timeout."""
    ser.timeout = timeout
    deadline = time.monotonic() + timeout

    # Scan for prefix byte
    while True:
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            return None
        ser.timeout = remaining
        b = ser.read(1)
        if not b:
            return None
        if b[0] == COMMAND_PREFIX:
            break

    remaining = deadline - time.monotonic()
    if remaining <= 0:
        return None
    ser.timeout = remaining
    hdr = ser.read(4)  # length, id, command, crc
    if len(hdr) < 4:
        return None
    length, resp_id, command, crc = hdr

    data = b""
    if length:
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            return None
        ser.timeout = remaining
        data = ser.read(length)
        if len(data) < length:
            return None

    expected_crc = calc_crc(length, command, data)
    return {
        "id": resp_id,
        "command": command,
        "length": length,
        "data": data,
        "crc": crc,
        "crc_ok": crc == expected_crc,
    }


def send_led(ser: serial.Serial):
    ser.write(build_led_packet())


# ---- Test modes -------------------------------------------------------------

def test_roundtrip(ser, delay, n, timeout):
    """Send n commands, waiting for a valid ACK after each before sending
    the next (with an *additional* `delay` seconds between send and next
    send, on top of whatever the round-trip itself costs). Returns
    (success_count, n, elapsed_seconds), where elapsed_seconds is the
    *measured* wall-clock time for the whole trial -- not derived from
    `delay`. This is what "rate" should be computed from, since at
    delay=0 the round-trip latency (write + firmware processing + read)
    is the real limiter, not some divide-by-zero fiction."""
    ser.reset_input_buffer()
    ok = 0
    t_start = time.monotonic()
    for i in range(n):
        send_led(ser)
        resp = read_response(ser, timeout)
        if resp and resp["crc_ok"] and resp["command"] == COMMAND_ACK:
            ok += 1
        if delay:
            time.sleep(delay)
    elapsed = time.monotonic() - t_start
    return ok, n, elapsed


def test_flood(ser, delay, n, timeout):
    """Send n commands back-to-back with only `delay` seconds between
    sends (no waiting for replies in between), then collect n replies
    afterward. Returns (success_count, n, elapsed_seconds).

    elapsed_seconds covers the whole trial (send phase + collecting all
    replies) as actually measured -- at delay=0 this still takes real,
    nonzero time (write() syscalls, USB bulk transfer framing, the
    firmware's IRQ handler + ring buffer + usb_thread all take time), so
    this is the number that should be used to compute cmd/s, never
    `1/delay`."""
    ser.reset_input_buffer()
    t_start = time.monotonic()
    for i in range(n):
        send_led(ser)
        if delay:
            time.sleep(delay)
    send_duration = time.monotonic() - t_start

    ok = 0
    # give it time proportional to what we sent, plus per-reply timeout
    end_deadline = time.monotonic() + timeout + send_duration
    for i in range(n):
        remaining = end_deadline - time.monotonic()
        if remaining <= 0:
            break
        resp = read_response(ser, remaining)
        if resp and resp["crc_ok"] and resp["command"] == COMMAND_ACK:
            ok += 1
        elif resp is None:
            break
    elapsed = time.monotonic() - t_start
    return ok, n, elapsed


def achieved_rate(ok, elapsed):
    """cmd/s actually measured, based on successful commands over real
    wall-clock time. Never derived from the requested delay."""
    if elapsed <= 0:
        return 0.0
    return ok / elapsed


# ---- Sweep / search logic ----------------------------------------------------

def find_max_rate(ser, mode, n, timeout, success_threshold, start_delay, min_delay, verbose):
    """Binary search the smallest stable inter-command delay that still
    achieves >= success_threshold success ratio. Returns
    (best_delay, measured_rate_at_best_delay).

    Rate is always the *measured* cmd/s from the trial, never 1/delay --
    at delay=0 the requested delay tells you nothing about the actual
    ceiling, which is set by write()/USB overhead and firmware processing
    time, not by our sleep() calls."""
    test_fn = test_roundtrip if mode == "roundtrip" else test_flood
    last_rate = 0.0

    def trial(delay):
        nonlocal last_rate
        ok, total, elapsed = test_fn(ser, delay, n, timeout)
        ratio = ok / total if total else 0.0
        rate = achieved_rate(ok, elapsed)
        last_rate = rate
        if verbose:
            print(f"  delay={delay*1000:8.3f} ms  measured_rate={rate:9.1f} cmd/s  "
                  f"success={ok}/{total} ({ratio*100:5.1f}%)  elapsed={elapsed*1000:.1f} ms")
        return ratio >= success_threshold

    lo, hi = min_delay, start_delay
    if not trial(hi):
        print(f"WARNING: even the slow starting delay ({hi*1000:.3f} ms) failed "
              f"the success threshold. Try increasing --start-delay.")
        return hi, last_rate
    hi_rate = last_rate

    if trial(lo):
        # Even the fastest requested delay was stable. That does NOT mean
        # the rate is infinite -- it means we've hit the real ceiling
        # (write()/USB/firmware), and last_rate is the measured number
        # for it. Report that instead of pretending it's unbounded.
        print(f"NOTE: even the fastest delay ({lo*1000:.3f} ms) passed, at a "
              f"measured {last_rate:.1f} cmd/s. That's likely the true ceiling "
              f"(write()/USB overhead + firmware processing), not an artifact "
              f"of --min-delay. Run --soak to confirm it holds over a longer run.")
        return lo, last_rate

    # Binary search between lo (fails) and hi (passes) for smallest passing delay
    best_rate = hi_rate
    for _ in range(20):
        mid = (lo + hi) / 2.0
        if trial(mid):
            hi = mid
            best_rate = last_rate
        else:
            lo = mid
        if hi - lo < 1e-5:  # 0.01 ms resolution
            break

    return hi, best_rate


def soak_test(ser, mode, delay, duration_s, timeout):
    """Run continuously at a fixed delay for `duration_s` seconds and
    report the measured sustained cmd/s and success ratio. Use this to
    confirm a delay=0 (or any) result actually holds up over time, since
    a short burst can pass while a longer run reveals ring-buffer
    overflow or drift."""
    test_fn = test_roundtrip if mode == "roundtrip" else test_flood
    t_start = time.monotonic()
    total_ok = 0
    total_n = 0
    # run in chunks so we can report progress and stop at duration_s
    chunk = 200
    while time.monotonic() - t_start < duration_s:
        ok, n, elapsed = test_fn(ser, delay, chunk, timeout)
        total_ok += ok
        total_n += n
        rate = achieved_rate(total_ok, time.monotonic() - t_start)
        print(f"  soak: {total_ok}/{total_n} ok so far, "
              f"sustained rate={rate:.1f} cmd/s, "
              f"t={time.monotonic() - t_start:5.1f}s / {duration_s}s")
    elapsed = time.monotonic() - t_start
    ratio = total_ok / total_n if total_n else 0.0
    rate = achieved_rate(total_ok, elapsed)
    return total_ok, total_n, ratio, rate


def sweep(ser, mode, n, timeout, delays, verbose=True):
    test_fn = test_roundtrip if mode == "roundtrip" else test_flood
    results = []
    print(f"\n{'delay (ms)':>12} {'measured cmd/s':>16} {'success':>10} {'ratio':>8}")
    print("-" * 52)
    for delay in delays:
        ok, total, elapsed = test_fn(ser, delay, n, timeout)
        ratio = ok / total if total else 0.0
        rate = achieved_rate(ok, elapsed)  # measured, never 1/delay
        print(f"{delay*1000:12.3f} {rate:16.1f} {ok:>4}/{total:<5} {ratio*100:7.1f}%")
        results.append((delay, ok, total, ratio, rate))
    return results


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                  formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--port", default="/dev/ttyACM0")
    ap.add_argument("--baud", type=int, default=115200,
                     help="Baud rate (often ignored by USB CDC-ACM, but set for compatibility)")
    ap.add_argument("--mode", choices=["roundtrip", "flood"], default="roundtrip",
                     help="roundtrip: wait for ACK after each send. "
                          "flood: send a burst, then collect replies (tests raw UART/ring-buffer throughput)")
    ap.add_argument("-n", type=int, default=200, help="Commands per trial")
    ap.add_argument("--timeout", type=float, default=0.5, help="Per-response read timeout (s)")
    ap.add_argument("--success-threshold", type=float, default=1.0,
                     help="Required success ratio to call a rate 'stable' (0-1)")
    ap.add_argument("--start-delay", type=float, default=0.02,
                     help="Slow starting inter-command delay in seconds for the search (known-good)")
    ap.add_argument("--min-delay", type=float, default=0.0,
                     help="Fastest inter-command delay to try, in seconds (0 = back-to-back)")
    ap.add_argument("--sweep", action="store_true",
                     help="Also print a full table sweeping delays geometrically "
                          "from --start-delay down to --min-delay-floor")
    ap.add_argument("--min-delay-floor", type=float, default=0.0005,
                     help="Smallest delay used for --sweep table (s)")
    ap.add_argument("--sweep-steps", type=int, default=12)
    ap.add_argument("--soak", type=float, default=0.0,
                     help="After the search, run this many seconds at the best "
                          "delay found to confirm the rate holds up over time "
                          "(recommended, especially when the best delay is 0)")
    args = ap.parse_args()

    print(f"Opening {args.port} @ {args.baud} baud, mode={args.mode}, n={args.n}/trial")
    ser = serial.Serial(args.port, args.baud, timeout=args.timeout)
    time.sleep(0.2)  # let the port settle
    ser.reset_input_buffer()
    ser.reset_output_buffer()

    try:
        if args.sweep:
            delays = []
            hi, lo = args.start_delay, args.min_delay_floor
            steps = max(args.sweep_steps, 1)
            for i in range(steps):
                frac = i / (steps - 1) if steps > 1 else 0
                # geometric interpolation from hi -> lo
                d = hi * ((lo / hi) ** frac) if hi > 0 else 0
                delays.append(d)
            sweep(ser, args.mode, args.n, args.timeout, delays)
            print()

        print("Binary-searching for max stable command rate...")
        best_delay, best_rate = find_max_rate(
            ser, args.mode, args.n, args.timeout,
            args.success_threshold, args.start_delay, args.min_delay,
            verbose=True,
        )

        print("\n=== Result ===")
        print(f"Max stable inter-command delay: {best_delay*1000:.3f} ms")
        print(f"Measured rate at that delay: {best_rate:.1f} commands/sec "
              f"(measured from actual elapsed time over {args.n} commands, "
              f"not derived from the delay)")
        print(f"(mode={args.mode}, n={args.n}, success_threshold={args.success_threshold*100:.0f}%)")

        if args.soak > 0:
            print(f"\nRunning {args.soak:.0f}s soak test at delay={best_delay*1000:.3f} ms "
                  f"to confirm this holds up over time...")
            ok, total, ratio, rate = soak_test(ser, args.mode, best_delay, args.soak, args.timeout)
            print(f"\nSoak result: {ok}/{total} ok ({ratio*100:.1f}%), "
                  f"sustained rate={rate:.1f} cmd/s over {args.soak:.0f}s")
            if ratio < args.success_threshold:
                print("WARNING: the short trial passed but the soak test did NOT hold up -- "
                      "the real stable rate is lower than reported above. Try a slower "
                      "--start-delay / larger --min-delay and re-run.")

    finally:
        ser.close()


if __name__ == "__main__":
    main()