musicmouse/python-backend/musicmouse/devices/wire.py

"""The serial wire protocol, as a pure codec.

No I/O and no asyncio here, so the exact byte layout the firmware depends on can be
pinned by ``tests/test_wire.py``. Everything must stay in sync with
``esp-firmware/src/Messages.h``; that file is the authority.

Frames in both directions are ``uint32 magic | uint8 type | uint16 payload_size``
followed by the payload, little-endian. The firmware also writes plain
``Serial.println`` log text on the same link, so the decoder resynchronises on
newlines and on the magic token.
"""

from __future__ import annotations

import struct
from dataclasses import dataclass, replace
from enum import IntEnum
from typing import TypeAlias

from musicmouse.effects import (
    EffectAlexaSwipeConfig,
    EffectCircularConfig,
    EffectRandomTwoColorInterpolationConfig,
    EffectReverseSwipe,
    EffectStaticConfig,
    EffectStaticDetailedConfig,
    EffectSwipeAndChange,
    LedEffect,
)
from musicmouse.events import (
    ButtonEvent,
    InputEvent,
    RfidTokenRead,
    RotaryTurned,
    TouchButtonPressed,
    TouchButtonReleased,
)
from musicmouse.hardware import (
    RFID_TAG_LENGTH,
    Button,
    ButtonAction,
    LedZone,
    RotaryDirection,
    TouchButton,
)

__all__ = [
    "MAGIC_FW_TO_HOST",
    "MAGIC_HOST_TO_FW",
    "Decoded",
    "FirmwareLog",
    "FrameDecoder",
    "HostCommand",
    "HostFrameDecoder",
    "MessageFwToHost",
    "MessageHostToFw",
    "ProtocolError",
    "SetButtonBrightness",
    "SetEffect",
    "UnsupportedEffectError",
    "encode_button_brightness",
    "encode_effect",
    "encode_input_event",
]

MAGIC_HOST_TO_FW = 0x1D6379E3
MAGIC_FW_TO_HOST = 0x10C65631

_HEADER = struct.Struct("<IBH")
_HEADER_SIZE = _HEADER.size  # 7
_MAGIC_FW_BYTES = struct.pack("<I", MAGIC_FW_TO_HOST)

#: Give up on resynchronising rather than buffering forever on a wedged link.
_MAX_BUFFER = 8192


class ProtocolError(Exception):
    """A frame arrived that could not be interpreted."""


class UnsupportedEffectError(Exception):
    """The firmware has no message for this effect on this LED zone."""


class MessageFwToHost(IntEnum):
    RFID_TOKEN_READ = 0
    ROTARY_ENCODER = 1
    TOUCH_BUTTON_PRESS = 2
    TOUCH_BUTTON_RELEASE = 3
    BUTTON_EVENT = 4


class MessageHostToFw(IntEnum):
    LED_WHEEL_EFFECT_STATIC = 0
    LED_WHEEL_EFFECT_ALEXA_SWIPE = 1
    LED_WHEEL_EFFECT_CIRCULAR = 2
    LED_WHEEL_EFFECT_RANDOM_TWO_COLOR_INTERPOLATION = 3
    LED_WHEEL_EFFECT_SWIPE_AND_CHANGE = 4
    LED_WHEEL_EFFECT_REVERSE_SWIPE = 5

    MOUSE_LED_EFFECT_STATIC = 6
    MOUSE_LED_EFFECT_CIRCULAR = 7
    MOUSE_LED_EFFECT_RANDOM_TWO_COLOR_INTERPOLATION = 8
    MOUSE_LED_EFFECT_SWIPE_AND_CHANGE = 9
    MOUSE_LED_EFFECT_REVERSE_SWIPE = 10

    SHELF_LED_EFFECT_STATIC = 15
    SHELF_LED_EFFECT_CIRCULAR = 16
    SHELF_LED_EFFECT_RANDOM_TWO_COLOR_INTERPOLATION = 17
    SHELF_LED_EFFECT_SWIPE_AND_CHANGE = 18
    SHELF_LED_EFFECT_REVERSE_SWIPE = 19
    SHELF_LED_EFFECT_STATIC_DETAILED = 20

    PREV_BUTTON_LED = 21
    NEXT_BUTTON_LED = 22


#: Which message id carries which effect, per zone. Note the asymmetry: only the ring
#: accepts AlexaSwipe on its own, and only the shelf accepts StaticDetailed.
_EFFECT_IDS: dict[LedZone, dict[type[LedEffect], MessageHostToFw]] = {
    LedZone.RING: {
        EffectStaticConfig: MessageHostToFw.LED_WHEEL_EFFECT_STATIC,
        EffectAlexaSwipeConfig: MessageHostToFw.LED_WHEEL_EFFECT_ALEXA_SWIPE,
        EffectCircularConfig: MessageHostToFw.LED_WHEEL_EFFECT_CIRCULAR,
        EffectRandomTwoColorInterpolationConfig: (
            MessageHostToFw.LED_WHEEL_EFFECT_RANDOM_TWO_COLOR_INTERPOLATION
        ),
        EffectSwipeAndChange: MessageHostToFw.LED_WHEEL_EFFECT_SWIPE_AND_CHANGE,
        EffectReverseSwipe: MessageHostToFw.LED_WHEEL_EFFECT_REVERSE_SWIPE,
    },
    LedZone.MOUSE: {
        EffectStaticConfig: MessageHostToFw.MOUSE_LED_EFFECT_STATIC,
        EffectCircularConfig: MessageHostToFw.MOUSE_LED_EFFECT_CIRCULAR,
        EffectRandomTwoColorInterpolationConfig: (
            MessageHostToFw.MOUSE_LED_EFFECT_RANDOM_TWO_COLOR_INTERPOLATION
        ),
        EffectSwipeAndChange: MessageHostToFw.MOUSE_LED_EFFECT_SWIPE_AND_CHANGE,
        EffectReverseSwipe: MessageHostToFw.MOUSE_LED_EFFECT_REVERSE_SWIPE,
    },
    LedZone.SHELF: {
        EffectStaticConfig: MessageHostToFw.SHELF_LED_EFFECT_STATIC,
        EffectCircularConfig: MessageHostToFw.SHELF_LED_EFFECT_CIRCULAR,
        EffectRandomTwoColorInterpolationConfig: (
            MessageHostToFw.SHELF_LED_EFFECT_RANDOM_TWO_COLOR_INTERPOLATION
        ),
        EffectSwipeAndChange: MessageHostToFw.SHELF_LED_EFFECT_SWIPE_AND_CHANGE,
        EffectReverseSwipe: MessageHostToFw.SHELF_LED_EFFECT_REVERSE_SWIPE,
        EffectStaticDetailedConfig: MessageHostToFw.SHELF_LED_EFFECT_STATIC_DETAILED,
    },
}

_BUTTON_LED_IDS: dict[Button, MessageHostToFw] = {
    Button.LEFT: MessageHostToFw.PREV_BUTTON_LED,
    Button.RIGHT: MessageHostToFw.NEXT_BUTTON_LED,
}


@dataclass(frozen=True, slots=True)
class FirmwareLog:
    """A ``Serial.println`` line from the firmware, interleaved with the frames."""

    text: str


Decoded: TypeAlias = InputEvent | FirmwareLog


# ------------------------------------------------------------------------- encoding


def _frame(message: MessageHostToFw, payload: bytes) -> bytes:
    return _HEADER.pack(MAGIC_HOST_TO_FW, message, len(payload)) + payload


def encode_effect(zone: LedZone, effect: LedEffect) -> bytes:
    """Encode ``effect`` as a frame for ``zone``.

    Raises:
        UnsupportedEffectError: if the firmware has no message for this combination.
    """
    try:
        message = _EFFECT_IDS[zone][type(effect)]
    except KeyError:
        supported = ", ".join(sorted(cls.__name__ for cls in _EFFECT_IDS[zone]))
        raise UnsupportedEffectError(
            f"{type(effect).__name__} cannot be sent to the {zone} LEDs "
            f"(supported there: {supported})"
        ) from None
    return _frame(message, effect.as_bytes())


def encode_button_brightness(button: Button, brightness: float) -> bytes:
    """Encode the backlight brightness (``0..1``) of the prev/next button."""
    if not 0 <= brightness <= 1:
        raise ValueError(f"brightness must be within 0..1, got {brightness}")
    try:
        message = _BUTTON_LED_IDS[button]
    except KeyError:
        raise ValueError(f"{button.slug} has no backlight") from None
    return _frame(message, struct.pack("<f", brightness))


def encode_input_event(event: InputEvent) -> bytes:
    """Encode an event as the firmware would send it.

    The inverse of the decoder, used by the simulator so that simulated hardware
    exercises the real codec rather than bypassing it.
    """
    match event:
        case RfidTokenRead(tag_id=tag_id):
            if len(tag_id) != RFID_TAG_LENGTH:
                raise ValueError(f"tag id must be {RFID_TAG_LENGTH} bytes, got {len(tag_id)}")
            message, payload = MessageFwToHost.RFID_TOKEN_READ, tag_id
        case RotaryTurned(position=position, increment=increment, direction=direction):
            message = MessageFwToHost.ROTARY_ENCODER
            payload = struct.pack("<iiB", position, increment, direction)
        case TouchButtonPressed(button=button):
            message, payload = MessageFwToHost.TOUCH_BUTTON_PRESS, struct.pack("<B", button)
        case TouchButtonReleased(button=button):
            message, payload = MessageFwToHost.TOUCH_BUTTON_RELEASE, struct.pack("<B", button)
        case ButtonEvent(button=push_button, action=action):
            message = MessageFwToHost.BUTTON_EVENT
            payload = struct.pack("<BB", push_button, action)
        case _:
            raise ValueError(f"{type(event).__name__} is not a firmware message")
    return struct.pack("<IBH", MAGIC_FW_TO_HOST, message, len(payload)) + payload


# ------------------------------------------------------------------------- decoding


def decode_message(msg_type: int, payload: bytes) -> InputEvent:
    """Turn one frame payload into an event.

    Raises:
        ProtocolError: on an unknown message type or a malformed payload.
    """
    try:
        message = MessageFwToHost(msg_type)
    except ValueError:
        raise ProtocolError(f"unknown message type {msg_type}") from None

    try:
        match message:
            case MessageFwToHost.RFID_TOKEN_READ:
                if len(payload) != RFID_TAG_LENGTH:
                    raise ProtocolError(
                        f"RFID payload must be {RFID_TAG_LENGTH} bytes, got {len(payload)}"
                    )
                return RfidTokenRead(tag_id=bytes(payload), source="device")
            case MessageFwToHost.ROTARY_ENCODER:
                position, increment, direction = struct.unpack("<iiB", payload)
                return RotaryTurned(
                    position=position,
                    increment=increment,
                    direction=RotaryDirection(direction),
                    source="device",
                )
            case MessageFwToHost.TOUCH_BUTTON_PRESS:
                return TouchButtonPressed(button=_touch_button(payload), source="device")
            case MessageFwToHost.TOUCH_BUTTON_RELEASE:
                return TouchButtonReleased(button=_touch_button(payload), source="device")
            case MessageFwToHost.BUTTON_EVENT:
                button_nr, event_nr = struct.unpack("<BB", payload)
                return ButtonEvent(
                    button=Button(button_nr), action=ButtonAction(event_nr), source="device"
                )
    except struct.error as exc:
        raise ProtocolError(f"malformed {message.name} payload {payload.hex()}: {exc}") from exc
    except ValueError as exc:
        raise ProtocolError(f"bad value in {message.name} payload {payload.hex()}: {exc}") from exc

    raise AssertionError(f"unhandled message {message}")  # pragma: no cover


def _touch_button(payload: bytes) -> TouchButton:
    if len(payload) != 1:
        raise ProtocolError(f"touch payload must be 1 byte, got {len(payload)}")
    return TouchButton(payload[0])


class FrameDecoder:
    """Incremental decoder for the byte stream coming from the firmware.

    Handles partial frames, several frames in one chunk, and interleaved log text.
    """

    def __init__(self) -> None:
        self._buffer = bytearray()

    def push(self, data: bytes) -> None:
        """Append freshly read bytes. Call :meth:`take` until it returns ``None``."""
        self._buffer += data

    @property
    def buffered(self) -> int:
        return len(self._buffer)

    def take(self) -> Decoded | None:
        """Return the next complete item, or ``None`` if more bytes are needed.

        Raises:
            ProtocolError: on a malformed frame. The offending frame has already been
                consumed, so the caller can log it and call ``take()`` again - which is
                why this is not a generator: an exception would close one for good.
        """
        buffer = self._buffer
        if not buffer:
            return None

        if buffer[:4] == _MAGIC_FW_BYTES:
            if len(buffer) < _HEADER_SIZE:
                return None
            _, msg_type, size = _HEADER.unpack_from(buffer)
            end = _HEADER_SIZE + size
            if len(buffer) < end:
                return None
            payload = bytes(buffer[_HEADER_SIZE:end])
            del buffer[:end]
            return decode_message(msg_type, payload)

        if len(buffer) < 4 and _MAGIC_FW_BYTES.startswith(buffer):
            return None  # could still become a frame header

        # Not a frame at offset 0, so it is firmware log text.
        if (newline := buffer.find(b"\n")) >= 0:
            line = bytes(buffer[:newline])
            del buffer[: newline + 1]
            return FirmwareLog(line.decode("utf-8", errors="replace").rstrip("\r"))

        # No newline yet: skip ahead to the next frame if one has already started.
        if (start := buffer.find(_MAGIC_FW_BYTES)) > 0:
            skipped = bytes(buffer[:start])
            del buffer[:start]
            return FirmwareLog(skipped.decode("utf-8", errors="replace").rstrip("\r"))

        if len(buffer) > _MAX_BUFFER:
            # Nothing recognisable and no end in sight - keep only what could still be
            # the beginning of a magic token straddling the next chunk.
            del buffer[: -len(_MAGIC_FW_BYTES) + 1]
        return None


def with_figure(event: RfidTokenRead, figure: str | None, *, known: bool) -> RfidTokenRead:
    """Attach a resolved figure name to a decoded tag read."""
    return replace(event, figure=figure, known=known)


# ------------------------------------------------- host -> firmware, read back
#
# Nothing in the running app needs this direction decoded - the firmware does that.
# The simulator uses it to report what the real device would have been told, which
# also means simulated hardware exercises the encoders rather than bypassing them.


@dataclass(frozen=True, slots=True)
class SetEffect:
    zone: LedZone
    effect: LedEffect

    def __repr__(self) -> str:
        return f"{self.zone} <- {self.effect}"


@dataclass(frozen=True, slots=True)
class SetButtonBrightness:
    button: Button
    brightness: float

    def __repr__(self) -> str:
        return f"{self.button.slug} backlight <- {self.brightness:.2f}"


HostCommand: TypeAlias = SetEffect | SetButtonBrightness

_ID_TO_EFFECT: dict[int, tuple[LedZone, type[LedEffect]]] = {
    message: (zone, effect_cls)
    for zone, effects in _EFFECT_IDS.items()
    for effect_cls, message in effects.items()
}
_ID_TO_BUTTON: dict[int, Button] = {
    message: button for button, message in _BUTTON_LED_IDS.items()
}
_MAGIC_HOST_BYTES = struct.pack("<I", MAGIC_HOST_TO_FW)


def decode_host_command(msg_type: int, payload: bytes) -> HostCommand:
    if (target := _ID_TO_EFFECT.get(msg_type)) is not None:
        zone, effect_cls = target
        try:
            return SetEffect(zone=zone, effect=effect_cls.from_bytes(payload))
        except (ValueError, struct.error) as exc:
            raise ProtocolError(f"malformed {effect_cls.__name__} payload: {exc}") from exc
    if (button := _ID_TO_BUTTON.get(msg_type)) is not None:
        try:
            (brightness,) = struct.unpack("<f", payload)
        except struct.error as exc:
            raise ProtocolError(f"malformed button brightness payload: {exc}") from exc
        return SetButtonBrightness(button=button, brightness=brightness)
    raise ProtocolError(f"unknown host-to-firmware message type {msg_type}")


class HostFrameDecoder:
    """Incremental decoder for the host -> firmware direction."""

    def __init__(self) -> None:
        self._buffer = bytearray()

    def push(self, data: bytes) -> None:
        self._buffer += data

    def take(self) -> HostCommand | None:
        """Next command, or ``None`` if more bytes are needed. See :meth:`FrameDecoder.take`."""
        if len(self._buffer) < _HEADER_SIZE:
            return None
        magic, msg_type, size = _HEADER.unpack_from(self._buffer)
        if magic != MAGIC_HOST_TO_FW:
            self._buffer.clear()
            raise ProtocolError(f"expected host-to-firmware magic, got {magic:#010x}")
        end = _HEADER_SIZE + size
        if len(self._buffer) < end:
            return None
        payload = bytes(self._buffer[_HEADER_SIZE:end])
        del self._buffer[:end]
        return decode_host_command(msg_type, payload)