#!/usr/bin/env python3
"""
Low level serial protocol for the sand table controller.

Speaks the framed protocol used by the firmware:

    struct command_message_t {
        uint8_t prefix;      // COMMAND_PREFIX (0x69)
        uint8_t length;      // number of valid bytes in `data`
        uint8_t id;          // COMMAND_ID (0x00)
        uint8_t command;     // commands_e
        uint8_t crc;         // 0x100 - (sum of every other byte in the message)
        uint8_t data[160];
    }

Every command the firmware accepts correctly is acknowledged with a
COMMAND_ACK message whose single data byte echoes the command that was
processed. For COMMAND_HOME and COMMAND_POLAR that ACK is only sent once the
steppers have physically finished moving, so waiting for it doubles as
"block until the move is done".

Position is now tracked entirely by the firmware (COMMAND_GET_POLAR), so
this module no longer keeps or persists any position of its own - callers
that need the current position call getPolar().
"""

import glob
import struct
import time
import serial


# -----------------------------------------------------------------------------
# Configuration
# -----------------------------------------------------------------------------

BAUDRATE = 115200

COMMAND_PREFIX = 0x69
COMMAND_ID = 0x00
COMMAND_DATA_SIZE = 160

# commands_e - order must match the firmware enum exactly
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_SET_OFFSET = 9
COMMAND_RESET_OFFSET = 10

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",
    COMMAND_SET_OFFSET: "SET_OFFSET",
    COMMAND_RESET_OFFSET: "RESET_OFFSET",
}

MIN_TRACK_SPEED = 1
MAX_TRACK_SPEED = 1000
MAX_BRIGHTNESS = 255

POLAR_STRUCT = struct.Struct("<ff")  # struct polar_t { float theta; float r; }

ser = None
PORT = None  # set once startSerial() successfully opens a port


# -----------------------------------------------------------------------------
# Connection
# -----------------------------------------------------------------------------

def _find_port():
    """Look for a /dev/ttyACMx device. Returns the first match (sorted), or
    None if nothing is plugged in."""
    candidates = sorted(glob.glob("/dev/ttyACM*"))
    if not candidates:
        return None
    if len(candidates) > 1:
        print(f"Multiple ttyACM ports found {candidates}, using {candidates[0]}")
    return candidates[0]


def startSerial(port=None):
    """Open the serial connection. If `port` isn't given, auto-detects a
    /dev/ttyACM* device."""
    global ser, PORT

    port = port or _find_port()
    if port is None:
        ser = None
        print("Failed to start serial connection: no /dev/ttyACM* device found")
        return

    try:
        ser = serial.Serial(port, BAUDRATE, timeout=1)
        PORT = port
        # Give the board a moment in case it resets on port open (common on
        # AVR / USB-CDC boards) before we start talking to it.
        time.sleep(2)
        ser.reset_input_buffer()
    except Exception as e:
        ser = None
        print(f"Failed to start serial connection on {port}: {e}")


def resetBuffer():
    if ser is not None:
        ser.reset_input_buffer()


def _ensure_serial():
    if ser is None:
        startSerial()
    return ser is not None


# -----------------------------------------------------------------------------
# Packet building / CRC (matches firmware's command_calculate_crc())
# -----------------------------------------------------------------------------

def _build_packet(command, data=b""):
    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])
    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


# -----------------------------------------------------------------------------
# Reading responses
# -----------------------------------------------------------------------------

def _read_message(timeout=2.0):
    """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 / no connection."""

    if not _ensure_serial():
        return None

    deadline = time.monotonic() + timeout
    old_timeout = ser.timeout

    try:
        while time.monotonic() < deadline:
            ser.timeout = max(0.01, deadline - time.monotonic())
            b = ser.read(1)
            if not b:
                continue
            if b[0] != COMMAND_PREFIX:
                continue  # resync: keep looking for the prefix byte

            header = ser.read(4)  # length, id, command, crc
            if len(header) < 4:
                continue
            length, msg_id, command, crc = header

            data = b""
            if length:
                data = ser.read(length)
                if len(data) < length:
                    continue  # malformed/short read, keep resyncing

            return {
                "prefix": b[0],
                "length": length,
                "id": msg_id,
                "command": command,
                "crc": crc,
                "data": data,
            }
    finally:
        ser.timeout = old_timeout

    return None


def _wait_for_ack(expected_command, timeout=10.0):
    """Wait for a COMMAND_ACK whose data[0] echoes expected_command. Returns
    the ACK payload (bytes) on success, or None on timeout/NACK."""

    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        remaining = deadline - time.monotonic()
        msg = _read_message(timeout=remaining)
        if msg is None:
            return None

        if msg["command"] == COMMAND_NACK:
            name = COMMAND_NAMES.get(expected_command, expected_command)
            print(f"  -> device NACKed command {name}")
            return None

        if msg["command"] == COMMAND_ACK:
            acked = msg["data"][0] if msg["data"] else None
            if acked == expected_command:
                return msg["data"]
            continue  # ACK for something else - keep waiting

    return None


# -----------------------------------------------------------------------------
# High level command helpers
# -----------------------------------------------------------------------------

def sendCommand(command, data=b"", wait_ack=True, ack_timeout=10.0):
    """Build, send, and (optionally) block for the completion ACK of a
    command. Returns the ACK payload (bytes, possibly empty) on success,
    True if wait_ack=False and the write succeeded, or None/False on
    failure."""

    if not _ensure_serial():
        return None if wait_ack else False

    packet = _build_packet(command, data)
    try:
        ser.write(packet)
        ser.flush()
    except Exception as e:
        print(f"Failed to write to serial: {e}")
        return None if wait_ack else False

    if not wait_ack:
        return True

    ack_data = _wait_for_ack(command, timeout=ack_timeout)
    if ack_data is None:
        name = COMMAND_NAMES.get(command, command)
        print(f"Timed out waiting for ACK on {name}")
    return ack_data


def sendPolar(theta, r, wait_ack=True, ack_timeout=15.0):
    """Move to (theta, r). Blocks until the steppers report the move is
    complete (unless wait_ack=False). Returns True on success, False/None
    on failure/timeout."""

    data = POLAR_STRUCT.pack(theta, r)
    result = sendCommand(COMMAND_POLAR, data, wait_ack=wait_ack, ack_timeout=ack_timeout)
    if not wait_ack:
        return result
    return result is not None


def sendPolarWithRetry(theta, r, retries=3, ack_timeout=15.0):
    """Same as sendPolar() but retries a few times on timeout/NACK before
    giving up (returns False)."""

    for attempt in range(retries):
        if sendPolar(theta, r, wait_ack=True, ack_timeout=ack_timeout):
            return True
        print(f"  retrying polar move to ({theta}, {r}) [{attempt + 1}/{retries}]")
        resetBuffer()
        time.sleep(0.05)
    return False


def getPolar(timeout=5.0):
    """Ask the device for its current (theta, r).

    ASSUMPTION: the reply is a command_message_t with command ==
    COMMAND_GET_POLAR and an 8-byte polar_t payload (same <ff> layout as
    what's sent to COMMAND_POLAR). I don't have the firmware's
    COMMAND_GET_POLAR case handler, so if the real reply looks different
    (e.g. comes back wrapped as an ACK instead) let me know and I'll adjust
    this.

    Returns (theta, r) or None on failure/timeout.
    """

    if not _ensure_serial():
        return None

    packet = _build_packet(COMMAND_GET_POLAR)
    try:
        ser.write(packet)
        ser.flush()
    except Exception as e:
        print(f"Failed to write to serial: {e}")
        return None

    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        msg = _read_message(timeout=deadline - time.monotonic())
        if msg is None:
            break
        if msg["command"] == COMMAND_NACK:
            print("  -> device NACKed GET_POLAR")
            return None
        if msg["command"] == COMMAND_GET_POLAR and len(msg["data"]) >= 8:
            return POLAR_STRUCT.unpack(msg["data"][:8])
        # anything else (e.g. a stray ACK for a previous command) - ignore and keep waiting

    print("Timed out waiting for GET_POLAR reply")
    return None


def home(wait_ack=True, ack_timeout=120.0):
    result = sendCommand(COMMAND_HOME, b"", wait_ack=wait_ack, ack_timeout=ack_timeout)
    if not wait_ack:
        return result
    return result is not None


def disableMotors(wait_ack=True, ack_timeout=5.0):
    result = sendCommand(COMMAND_DISABLE_MOTORS, b"", wait_ack=wait_ack, ack_timeout=ack_timeout)
    if not wait_ack:
        return result
    return result is not None


def setMotorSpeed(speed, wait_ack=True, ack_timeout=5.0):
    """ASSUMPTION: motor speed is encoded as a uint16 little-endian value,
    matching the LED track speed convention. I don't have the firmware's
    COMMAND_MOTOR_SPEED handler, so confirm this - if it actually expects a
    float or a different width, this needs to change."""

    speed = max(0, min(int(speed), 0xFFFF))
    data = struct.pack("<H", speed)
    result = sendCommand(COMMAND_MOTOR_SPEED, data, wait_ack=wait_ack, ack_timeout=ack_timeout)
    if not wait_ack:
        return result
    return result is not None


def setOffset(m1_delta, m2_delta, wait_ack=True, ack_timeout=5.0):
    """Add to the firmware's stored per-motor calibration offset (steps).
    This is ADDITIVE (stepper_add_offset), not absolute - each call nudges
    the offset further, matching the old "+1/-1/+10/-10" calibrate button
    behavior. Both deltas are signed 8-bit (-128..127)."""

    m1 = max(-128, min(int(m1_delta), 127))
    m2 = max(-128, min(int(m2_delta), 127))
    data = struct.pack("<bb", m1, m2)
    result = sendCommand(COMMAND_SET_OFFSET, data, wait_ack=wait_ack, ack_timeout=ack_timeout)
    if not wait_ack:
        return result
    return result is not None


def resetOffset(wait_ack=True, ack_timeout=5.0):
    """Zero out the firmware's stored calibration offset entirely."""
    result = sendCommand(COMMAND_RESET_OFFSET, b"", wait_ack=wait_ack, ack_timeout=ack_timeout)
    if not wait_ack:
        return result
    return result is not None


def setLedTrack(colors, speed, wait_ack=True, ack_timeout=5.0):
    """colors: list of (r, g, b, w) tuples, 0-255 each. speed: track speed
    (uint16). A single solid color is just a one-element list. The device
    holds no LED track state of its own, so the *full* track has to be
    re-sent on every change - matches the reference LED script's
    "speed + count + colors" payload layout."""

    if not colors:
        raise ValueError("colors must contain at least one (r, g, b, w) tuple")
    # 3 header bytes (speed lo/hi + count) + 4 bytes/color must fit in 160
    if 3 + len(colors) * 4 > COMMAND_DATA_SIZE:
        raise ValueError("too many colors for a single COMMAND_LED payload")

    payload = bytearray()
    speed = max(0, min(int(speed), 0xFFFF))
    payload.append(speed & 0xFF)
    payload.append((speed >> 8) & 0xFF)
    payload.append(len(colors))
    for color in colors:
        payload += struct.pack("BBBB", *(max(0, min(int(c), 255)) for c in color))

    result = sendCommand(COMMAND_LED, bytes(payload), wait_ack=wait_ack, ack_timeout=ack_timeout)
    if not wait_ack:
        return result
    return result is not None


startSerial()