From 2604ff274bd452c55848c93f9bb117b4fc324407 Mon Sep 17 00:00:00 2001
From: Joey Yakimowich-Payne <jyapayne@pm.me>
Date: Mon, 16 Mar 2026 11:50:25 -0600
Subject: [PATCH] feat(uart): add v2 protocol with IMU sample support

- Add UART_PROTOCOL_VERSION=2, IMUSample dataclass, conversion constants
- Rewrite SwitchReport.to_bytes() for versioned v2 framing with checksum
- Add ACCEL_LSB_PER_G=4096, GYRO_LSB_PER_RAD_S=818.5, MS2_PER_G constants
- Add SENSOR_ACCEL/SENSOR_GYRO SDL type constants with fallback
- Add protocol round-trip tests (8 tests, all passing)
---
 src/switch_pico_bridge/switch_pico_uart.py |  97 ++++++++++++++---
 tests/__init__.py                          |   0
 tests/test_uart_protocol.py                | 120 +++++++++++++++++++++
 3 files changed, 204 insertions(+), 13 deletions(-)
 create mode 100644 tests/__init__.py
 create mode 100644 tests/test_uart_protocol.py

diff --git a/src/switch_pico_bridge/switch_pico_uart.py b/src/switch_pico_bridge/switch_pico_uart.py
index eb4fa25..2fdab53 100644
--- a/src/switch_pico_bridge/switch_pico_uart.py
+++ b/src/switch_pico_bridge/switch_pico_uart.py
@@ -12,6 +12,7 @@ depending on SDL. It mirrors the framing in ``switch-pico.cpp``:
 
 from __future__ import annotations
 
+import math
 import struct
 import time
 import threading
@@ -23,9 +24,28 @@ import serial
 from serial.tools import list_ports, list_ports_common
 
 UART_HEADER = 0xAA
+UART_PROTOCOL_VERSION = 0x02
 RUMBLE_HEADER = 0xBB
 RUMBLE_TYPE_RUMBLE = 0x01
 UART_BAUD = 921600
+IMU_SAMPLES_PER_REPORT = 3
+
+MS2_PER_G = 9.80665
+RAD_TO_DEG = 180.0 / math.pi
+ACCEL_LSB_PER_G = 4096.0
+GYRO_LSB_PER_RAD_S = 818.5
+
+try:
+    import sdl2 as _sdl2  # type: ignore[import-not-found]
+
+    _sensor_accel = getattr(_sdl2, "SDL_SENSOR_ACCEL", 1)
+    _sensor_gyro = getattr(_sdl2, "SDL_SENSOR_GYRO", 2)
+except ImportError:
+    _sensor_accel = 1
+    _sensor_gyro = 2
+
+SENSOR_ACCEL: int = _sensor_accel
+SENSOR_GYRO: int = _sensor_gyro
 
 
 class SwitchButton(IntFlag):
@@ -62,9 +82,9 @@ def _is_usb_serial_path(path: str) -> bool:
     """Heuristic for USB serial path prefixes."""
     lower = path.lower()
     usb_prefixes = (
-        "/dev/ttyusb",   # Linux USB serial
-        "/dev/ttyacm",   # Linux CDC ACM
-        "/dev/cu.usb",   # macOS cu/tty USB adapters
+        "/dev/ttyusb",  # Linux USB serial
+        "/dev/ttyacm",  # Linux CDC ACM
+        "/dev/cu.usb",  # macOS cu/tty USB adapters
         "/dev/tty.usb",
     )
     if lower.startswith(usb_prefixes):
@@ -144,6 +164,7 @@ def first_serial_port(
         return None
     return ports[0]["device"]
 
+
 def clamp_byte(value: Union[int, float]) -> int:
     """Clamp a numeric value to the 0-255 byte range."""
     return max(0, min(255, int(value)))
@@ -202,6 +223,21 @@ def str_to_dpad(flags: Mapping[str, bool]) -> SwitchDpad:
     return SwitchDpad.CENTER
 
 
+def compute_checksum(data: bytes) -> int:
+    """Compute UART checksum as sum of bytes modulo 256."""
+    return sum(data) & 0xFF
+
+
+@dataclass
+class IMUSample:
+    accel_x: int = 0
+    accel_y: int = 0
+    accel_z: int = 0
+    gyro_x: int = 0
+    gyro_y: int = 0
+    gyro_z: int = 0
+
+
 @dataclass
 class SwitchReport:
     buttons: int = 0
@@ -210,13 +246,38 @@ class SwitchReport:
     ly: int = 128
     rx: int = 128
     ry: int = 128
+    imu_samples: List[IMUSample] = field(default_factory=list)
 
     def to_bytes(self) -> bytes:
-        """Serialize the report into the UART packet format."""
-        return struct.pack(
-            "<BHBBBBB", UART_HEADER, self.buttons & 0xFFFF, self.hat & 0xFF, self.lx, self.ly, self.rx, self.ry
+        """Serialize the report into UART v2 framed packet format."""
+        count = min(len(self.imu_samples), IMU_SAMPLES_PER_REPORT)
+        payload = struct.pack(
+            "<HBBBBBB",
+            self.buttons & 0xFFFF,
+            int(self.hat) & 0xFF,
+            clamp_byte(self.lx),
+            clamp_byte(self.ly),
+            clamp_byte(self.rx),
+            clamp_byte(self.ry),
+            count,
         )
 
+        for i in range(count):
+            sample = self.imu_samples[i]
+            payload += struct.pack(
+                "<hhhhhh",
+                max(-32768, min(32767, int(sample.accel_x))),
+                max(-32768, min(32767, int(sample.accel_y))),
+                max(-32768, min(32767, int(sample.accel_z))),
+                max(-32768, min(32767, int(sample.gyro_x))),
+                max(-32768, min(32767, int(sample.gyro_y))),
+                max(-32768, min(32767, int(sample.gyro_z))),
+            )
+
+        payload_len = len(payload)
+        frame = bytes([UART_HEADER, UART_PROTOCOL_VERSION, payload_len]) + payload
+        return frame + bytes([compute_checksum(frame)])
+
 
 class PicoUART:
     def __init__(self, port: str, baudrate: int = UART_BAUD) -> None:
@@ -267,15 +328,15 @@ class PicoUART:
                     del self._buffer[:start]
                 return None
 
-            frame = self._buffer[start:start + 11]
-            checksum = sum(frame[:10]) & 0xFF
+            frame = self._buffer[start : start + 11]
+            checksum = compute_checksum(bytes(frame[:10]))
 
             if frame[1] == RUMBLE_TYPE_RUMBLE and checksum == frame[10]:
                 payload = bytes(frame[2:10])
-                del self._buffer[:start + 11]
+                del self._buffer[: start + 11]
                 return payload
 
-            del self._buffer[:start + 1]
+            del self._buffer[: start + 1]
 
     def close(self) -> None:
         """Close the UART connection."""
@@ -434,14 +495,20 @@ class SwitchUARTClient:
         self.state.move_right_stick(x, y)
         self.send()
 
-    def press_for(self, duration: float, *buttons: SwitchButton | SwitchDpad | int) -> None:
+    def press_for(
+        self, duration: float, *buttons: SwitchButton | SwitchDpad | int
+    ) -> None:
         """Press buttons/hat for a duration, then release."""
         self.press(*buttons)
         time.sleep(max(0.0, duration))
         self.release(*buttons)
 
     def move_left_stick_for(
-        self, x: Union[int, float], y: Union[int, float], duration: float, neutral_after: bool = True
+        self,
+        x: Union[int, float],
+        y: Union[int, float],
+        duration: float,
+        neutral_after: bool = True,
     ) -> None:
         """Move left stick for a duration, optionally returning it to neutral afterward."""
         self.move_left_stick(x, y)
@@ -451,7 +518,11 @@ class SwitchUARTClient:
             self.send()
 
     def move_right_stick_for(
-        self, x: Union[int, float], y: Union[int, float], duration: float, neutral_after: bool = True
+        self,
+        x: Union[int, float],
+        y: Union[int, float],
+        duration: float,
+        neutral_after: bool = True,
     ) -> None:
         """Move right stick for a duration, optionally returning it to neutral afterward."""
         self.move_right_stick(x, y)
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/test_uart_protocol.py b/tests/test_uart_protocol.py
new file mode 100644
index 0000000..be3f8bd
--- /dev/null
+++ b/tests/test_uart_protocol.py
@@ -0,0 +1,120 @@
+"""Tests for UART v2 protocol serialization in switch_pico_uart."""
+
+import struct
+import pytest
+from switch_pico_bridge.switch_pico_uart import (
+    SwitchReport,
+    IMUSample,
+    SwitchDpad,
+    UART_HEADER,
+    UART_PROTOCOL_VERSION,
+    ACCEL_LSB_PER_G,
+    GYRO_LSB_PER_RAD_S,
+    MS2_PER_G,
+    compute_checksum,
+)
+
+
+def test_v2_frame_with_imu_samples():
+    """V2 frame with 3 IMU samples should be 48 bytes with correct layout."""
+    r = SwitchReport(
+        buttons=0,
+        imu_samples=[
+            IMUSample(100, -200, 4096, 50, -50, 0),
+            IMUSample(101, -201, 4097, 51, -51, 1),
+            IMUSample(102, -202, 4098, 52, -52, 2),
+        ],
+    )
+    data = r.to_bytes()
+    assert len(data) == 48, f"Expected 48 bytes, got {len(data)}"
+    assert data[0] == UART_HEADER  # 0xAA
+    assert data[1] == UART_PROTOCOL_VERSION  # 0x02
+    assert data[2] == 44  # payload_len
+    assert data[10] == 3  # imu_count
+    # Verify checksum
+    assert data[-1] == compute_checksum(data[:-1])
+    # Verify first sample accel_x (int16 LE at byte 11)
+    ax0 = struct.unpack_from("<h", data, 11)[0]
+    assert ax0 == 100, f"Expected accel_x=100, got {ax0}"
+    # Verify first sample gyro_z (int16 LE at bytes 21-22)
+    gz0 = struct.unpack_from("<h", data, 21)[0]
+    assert gz0 == 0, f"Expected gyro_z=0, got {gz0}"
+
+
+def test_v2_frame_no_imu():
+    """V2 frame with no IMU samples should be 12 bytes."""
+    r = SwitchReport(
+        buttons=0x0004, hat=SwitchDpad.CENTER, lx=128, ly=128, rx=128, ry=128
+    )
+    data = r.to_bytes()
+    assert len(data) == 12, f"Expected 12 bytes, got {len(data)}"
+    assert data[0] == UART_HEADER
+    assert data[1] == UART_PROTOCOL_VERSION
+    assert data[2] == 8  # payload_len
+    assert data[10] == 0  # imu_count
+    assert data[-1] == compute_checksum(data[:-1])
+
+
+def test_checksum_validation():
+    """Checksum should match sum of all preceding bytes & 0xFF."""
+    r = SwitchReport(buttons=0x0001)
+    data = r.to_bytes()
+    expected_checksum = sum(data[:-1]) & 0xFF
+    assert data[-1] == expected_checksum
+    # Corrupt a byte and verify mismatch
+    corrupted = bytearray(data)
+    corrupted[3] ^= 0xFF  # flip bits in first payload byte
+    recalculated = sum(corrupted[:-1]) & 0xFF
+    assert corrupted[-1] != recalculated, "Checksum should not match corrupted data"
+
+
+def test_accel_scale_gravity():
+    """1G (9.80665 m/s²) should convert to ~4096 raw counts."""
+    # convert_accel_to_raw(9.80665) ≈ 4096
+    raw = int(round((MS2_PER_G / MS2_PER_G) * ACCEL_LSB_PER_G))
+    assert abs(raw - 4096) <= 5, f"Expected ~4096 for 1G, got {raw}"
+
+
+def test_gyro_scale_one_rad():
+    """1.0 rad/s should convert to ~818 raw counts."""
+    raw = int(round(1.0 * GYRO_LSB_PER_RAD_S))
+    assert abs(raw - 818) <= 5, f"Expected ~818 for 1 rad/s, got {raw}"
+
+
+def test_imu_sample_dataclass():
+    """IMUSample fields accept int16 range values."""
+    s = IMUSample(
+        accel_x=32767, accel_y=-32768, accel_z=0, gyro_x=100, gyro_y=-100, gyro_z=1000
+    )
+    assert s.accel_x == 32767
+    assert s.accel_y == -32768
+    assert s.gyro_z == 1000
+    # Values outside int16 range are clamped in to_bytes()
+    s2 = IMUSample(accel_x=99999)
+    r = SwitchReport(imu_samples=[s2])
+    data = r.to_bytes()
+    ax = struct.unpack_from("<h", data, 11)[0]
+    assert ax == 32767, f"Expected clamped value 32767, got {ax}"
+
+
+def test_backward_compat_switch_report():
+    """SwitchReport with no imu_samples produces valid v2 frame (backward compat)."""
+    r = SwitchReport(buttons=0x000A, lx=200, ly=50, rx=128, ry=128)
+    data = r.to_bytes()
+    assert len(data) == 12
+    assert data[1] == 0x02  # still v2
+    # Buttons at bytes 3-4
+    buttons = struct.unpack_from("<H", data, 3)[0]
+    assert buttons == 0x000A
+    # lx at byte 6
+    assert data[6] == 200
+
+
+def test_max_imu_samples_capped():
+    """Providing >3 IMU samples should cap at 3."""
+    samples = [IMUSample(i, 0, 0, 0, 0, 0) for i in range(5)]
+    r = SwitchReport(imu_samples=samples)
+    data = r.to_bytes()
+    assert len(data) == 48  # 3 samples, not 5
+    assert data[10] == 3
+    assert data[2] == 44  # payload_len for 3 samples