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

Reads a "thetarho" coordinate file (the format sandify exports, e.g.:

    #
    # File name: 'sandify'
    # File type: thetarho
    #

    # BEGIN LAYER: 0 Polygon
    0.00000 0.00000
    1.57080 0.00800
    ...

) and streams each (theta, rho) point to the sand-table controller over
serial as a COMMAND_POLAR message, matching the firmware's
`struct command_message_t` / `struct polar_t` protocol.

Usage examples:

    # Just stream the file, waiting for the completion ACK after every point
    ./send_thetarho.py pattern.thr

    # Home first, then stream, don't wait for ACKs (fire-and-forget)
    ./send_thetarho.py pattern.thr --home --no-wait-ack

    # Different port / baud
    ./send_thetarho.py pattern.thr -p /dev/ttyACM0 -b 115200
"""

import argparse
import struct
import sys
import time

import serial

# -----------------------------------------------------------------------------
# Protocol constants (must match firmware)
# -----------------------------------------------------------------------------

COMMAND_PREFIX = 0x69
COMMAND_ID = 0x00

COMMAND_DATA_SIZE = 160  # sizeof(msg->data) on the firmware side

# commands_e
COMMAND_ACK = 0
COMMAND_NACK = 1
COMMAND_LED = 2
COMMAND_HOME = 3
COMMAND_DISABLE_MOTORS = 4
COMMAND_MOTOR_STEP = 5
COMMAND_MOTOR_SPEED = 6
COMMAND_POLAR = 7
COMMAND_GET_POLAR = 8

COMMAND_NAMES = {
    COMMAND_ACK: "ACK",
    COMMAND_NACK: "NACK",
    COMMAND_LED: "LED",
    COMMAND_HOME: "HOME",
    COMMAND_DISABLE_MOTORS: "DISABLE_MOTORS",
    COMMAND_MOTOR_STEP: "MOTOR_STEP",
    COMMAND_MOTOR_SPEED: "MOTOR_SPEED",
    COMMAND_POLAR: "POLAR",
    COMMAND_GET_POLAR: "GET_POLAR",
}

HEADER_LEN = 5  # prefix, length, id, command, crc

# struct polar_t { float theta; float r; };  -> two 32-bit floats
POLAR_STRUCT = struct.Struct("<ff")


# -----------------------------------------------------------------------------
# Packet building / CRC (mirrors command_calculate_crc() on the firmware)
# -----------------------------------------------------------------------------

def build_packet(command: int, data: bytes = b"") -> bytearray:
    """Build a command_message_t packet. `data` is only the bytes actually
    used (the command does NOT need to fill the full 160-byte buffer)."""

    if len(data) > COMMAND_DATA_SIZE:
        raise ValueError(
            f"data length {len(data)} exceeds COMMAND_DATA_SIZE ({COMMAND_DATA_SIZE})"
        )

    length = len(data)

    packet = bytearray(
        [
            COMMAND_PREFIX,
            length,
            COMMAND_ID,
            command,
            0x00,  # crc placeholder
        ]
    )
    packet += data

    crc_sum = 0
    for i, byte in enumerate(packet):
        if i == 4:  # skip the crc field itself
            continue
        crc_sum += byte

    packet[4] = (0x100 - (crc_sum & 0xFF)) & 0xFF

    return packet


def polar_packet(theta: float, r: float) -> bytearray:
    return build_packet(COMMAND_POLAR, POLAR_STRUCT.pack(theta, r))


# -----------------------------------------------------------------------------
# Reading responses back from the device
# -----------------------------------------------------------------------------

def read_message(ser: serial.Serial, timeout: float = 2.0, debug: bool = False):
    """Read one command_message_t from the serial port, resyncing on the
    0x69 prefix byte. Returns a dict {prefix, length, id, command, crc, data}
    or None on timeout."""

    deadline = time.monotonic() + timeout
    old_timeout = ser.timeout
    discarded = bytearray()
    raw_rx = bytearray()  # every byte read during this call, success or not

    def flush_discarded():
        if debug and discarded:
            print(f"  [debug] discarded while resyncing: {bytes(discarded).hex(' ')}")
            discarded.clear()

    try:
        while time.monotonic() < deadline:
            ser.timeout = max(0.01, deadline - time.monotonic())
            b = ser.read(1)
            if not b:
                continue
            raw_rx += b
            if debug:
                print(f"  [debug] RX byte: {b.hex()}")
            if b[0] != COMMAND_PREFIX:
                discarded += b  # resync: keep looking for the prefix byte
                continue

            header = ser.read(4)  # length, id, command, crc
            raw_rx += header
            if debug and header:
                print(f"  [debug] RX header: {header.hex(' ')}")
            if len(header) < 4:
                discarded += b + header
                continue
            length, msg_id, command, crc = header

            data = b""
            if length:
                data = ser.read(length)
                raw_rx += data
                if debug:
                    print(f"  [debug] RX data ({len(data)}/{length}): {data.hex(' ')}")
                if len(data) < length:
                    discarded += b + header + data
                    continue  # malformed/short read, keep resyncing

            flush_discarded()

            msg = {
                "prefix": b[0],
                "length": length,
                "id": msg_id,
                "command": command,
                "crc": crc,
                "data": data,
            }

            if debug:
                full = bytearray([b[0], length, msg_id, command, crc]) + bytearray(data)
                name = COMMAND_NAMES.get(command, hex(command))
                print(f"  [debug] recv {name} <- {bytes(full).hex(' ')}")

            return msg
    finally:
        flush_discarded()
        if debug and not raw_rx:
            print("  [debug] RX: nothing received before timeout")
        elif debug:
            print(f"  [debug] RX total this call: {bytes(raw_rx).hex(' ')}")
        ser.timeout = old_timeout

    return None


def wait_for_ack(ser: serial.Serial, expected_command: int, timeout: float = 10.0, debug: bool = False) -> bool:
    """Wait for an ACK/NACK message whose data[0] echoes expected_command
    (this is how the firmware signals a HOME/POLAR move has completed)."""

    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        remaining = deadline - time.monotonic()
        msg = read_message(ser, timeout=remaining, debug=debug)
        if msg is None:
            return False

        if msg["command"] == COMMAND_NACK:
            print("  -> received NACK from device")
            return False

        if msg["command"] == COMMAND_ACK:
            acked = msg["data"][0] if msg["data"] else None
            if acked == expected_command:
                return True
            # ACK for something else (e.g. a stray immediate ACK) - keep waiting
            continue

    return False


# -----------------------------------------------------------------------------
# Thetarho file parsing
# -----------------------------------------------------------------------------

def parse_thetarho(path: str):
    """Yield (theta, r) tuples from a sandify-style thetarho file, skipping
    comments (#...) and blank lines."""

    points = []
    with open(path, "r") as f:
        for lineno, raw_line in enumerate(f, start=1):
            line = raw_line.strip()
            if not line or line.startswith("#"):
                continue

            parts = line.split()
            if len(parts) != 2:
                print(f"warning: skipping malformed line {lineno}: {raw_line!r}", file=sys.stderr)
                continue

            try:
                theta = float(parts[0])
                r = float(parts[1])
            except ValueError:
                print(f"warning: skipping non-numeric line {lineno}: {raw_line!r}", file=sys.stderr)
                continue

            points.append((theta, r))

    return points


# -----------------------------------------------------------------------------
# Main
# -----------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(description="Stream a thetarho file to the sand table over serial.")
    parser.add_argument("file", help="Path to the thetarho (.thr) file")
    parser.add_argument("-p", "--port", default="/dev/ttyACM1", help="Serial port (default: %(default)s)")
    parser.add_argument("-b", "--baud", type=int, default=115200, help="Baud rate (default: %(default)s)")
    parser.add_argument("--home", action="store_true", help="Send a HOME command before streaming points")
    parser.add_argument(
        "--no-wait-ack",
        dest="wait_ack",
        action="store_false",
        default=True,
        help="Don't wait for the completion ACK between points (fire-and-forget)",
    )
    parser.add_argument(
        "--ack-timeout",
        type=float,
        default=15.0,
        help="Seconds to wait for a completion ACK before giving up on a point (default: %(default)s)",
    )
    parser.add_argument(
        "--delay",
        type=float,
        default=0.0,
        help="Extra delay in seconds between points when not waiting for ACKs",
    )
    parser.add_argument("-v", "--verbose", action="store_true", help="Print each packet sent")
    parser.add_argument(
        "--debug",
        action="store_true",
        help="Print every sent packet's bytearray, every parsed response, and any "
        "stray bytes discarded while resyncing on the serial line. Implies --verbose.",
    )
    args = parser.parse_args()
    if args.debug:
        args.verbose = True

    points = parse_thetarho(args.file)
    if not points:
        print(f"No valid points found in {args.file}", file=sys.stderr)
        sys.exit(1)

    print(f"Loaded {len(points)} points from {args.file}")

    with serial.Serial(args.port, args.baud, timeout=1) as ser:
        # Give the device a moment in case it resets on port open (common on
        # AVR/USB-CDC boards).
        time.sleep(2)
        ser.reset_input_buffer()

        if args.home:
            print("Homing...")
            packet = build_packet(COMMAND_HOME)
            if args.verbose:
                print("  ->", bytes(packet).hex(" "))
            ser.write(packet)
            ser.flush()
            if args.wait_ack:
                if not wait_for_ack(ser, COMMAND_HOME, timeout=args.ack_timeout, debug=args.debug):
                    print("Timed out / failed waiting for HOME to complete", file=sys.stderr)
                    sys.exit(1)
            print("Homed.")

        for i, (theta, r) in enumerate(points):
            packet = polar_packet(theta, r)
            if args.verbose:
                print(f"[{i + 1}/{len(points)}] theta={theta:.5f} r={r:.5f} -> {bytes(packet).hex(' ')}")
            else:
                print(f"\r[{i + 1}/{len(points)}] theta={theta:.5f} r={r:.5f}", end="", flush=True)

            ser.write(packet)
            ser.flush()

            if args.wait_ack:
                if not wait_for_ack(ser, COMMAND_POLAR, timeout=args.ack_timeout, debug=args.debug):
                    print(f"\nTimed out / failed waiting for point {i + 1} to complete", file=sys.stderr)
                    sys.exit(1)
            elif args.delay:
                time.sleep(args.delay)

        if not args.verbose:
            print()

    print("Done.")


if __name__ == "__main__":
    main()