led2/lightsync/protocol/simulator.py

"""Software UDP simulator — receives and logs WLED-compatible packets.

Used for hardware-free protocol testing. Listens on a UDP port and parses
all incoming packets (DRGB, DRGBW, DNRGB, animation commands), logging
them for test assertions.

Usage as standalone server:
    python -m lightsync.protocol.simulator

Usage in tests:
    loop = asyncio.get_running_loop()
    transport, protocol = await loop.create_datagram_endpoint(
        SimulatorProtocol, local_addr=("127.0.0.1", 0)
    )
    port = transport.get_extra_info("sockname")[1]
    # ... send packets to 127.0.0.1:port ...
    assert protocol.packets[0].protocol == "DRGB"
    transport.close()
"""
import asyncio
import struct
from dataclasses import dataclass, field
from lightsync.protocol.drgb import PROTOCOL_DRGB, PROTOCOL_DRGBW, PROTOCOL_DNRGB, WLED_PORT
from lightsync.protocol.animation_cmd import ANIM_CMD_MARKER, decode_animation_cmd

PROTOCOL_NAMES = {
    PROTOCOL_DRGB: "DRGB",
    PROTOCOL_DRGBW: "DRGBW",
    PROTOCOL_DNRGB: "DNRGB",
    ANIM_CMD_MARKER: "ANIM_CMD",
}


@dataclass
class PacketLog:
    """Parsed representation of a received UDP packet."""
    protocol: str
    addr: tuple
    led_count: int = 0
    start_index: int = 0
    raw: bytes = b""
    decoded: dict | None = None


class SimulatorProtocol(asyncio.DatagramProtocol):
    """Async UDP protocol that parses and logs incoming WLED-compatible packets.

    Supports all four packet types:
    - DRGB (0x02): RGB pixel data
    - DRGBW (0x03): RGBW pixel data (SK6812)
    - DNRGB (0x04): Multi-packet RGB with start index
    - ANIM_CMD (0xAC): LightSync animation command
    """

    def __init__(self) -> None:
        self.packets: list[PacketLog] = []
        self.transport: asyncio.DatagramTransport | None = None

    def connection_made(self, transport: asyncio.DatagramTransport) -> None:
        self.transport = transport

    def datagram_received(self, data: bytes, addr: tuple) -> None:
        if len(data) < 2:
            return

        proto_type = data[0]
        name = PROTOCOL_NAMES.get(proto_type, f"UNKNOWN(0x{proto_type:02X})")

        log = PacketLog(protocol=name, addr=addr, raw=data)

        if proto_type == PROTOCOL_DRGB:
            # Format: [0x02][timeout][R][G][B]...  — 3 bytes per LED after 2-byte header
            log.led_count = (len(data) - 2) // 3

        elif proto_type == PROTOCOL_DRGBW:
            # Format: [0x03][timeout][R][G][B][W]...  — 4 bytes per LED after 2-byte header
            log.led_count = (len(data) - 2) // 4

        elif proto_type == PROTOCOL_DNRGB:
            # Format: [0x04][timeout][start_hi][start_lo][R][G][B]...
            # Start index is big-endian uint16 at bytes 2-3; pixels start at byte 4
            if len(data) >= 4:
                log.start_index = struct.unpack_from(">H", data, 2)[0]
                log.led_count = (len(data) - 4) // 3

        elif proto_type == ANIM_CMD_MARKER:
            try:
                log.decoded = decode_animation_cmd(data)
            except Exception:
                log.decoded = None

        self.packets.append(log)

        # Human-readable log for standalone mode
        details = f"{log.led_count} LEDs" if log.led_count else ""
        if proto_type == PROTOCOL_DNRGB:
            details += f" start={log.start_index}"
        if log.decoded:
            details += f" cmd={log.decoded}"
        print(f"[SIM] {name} from {addr}: {details or len(data)} bytes")

    def error_received(self, exc: Exception) -> None:
        print(f"[SIM] Error: {exc}")

    def connection_lost(self, exc: Exception | None) -> None:
        pass


async def run_simulator(
    host: str = "0.0.0.0",
    port: int = WLED_PORT,
) -> tuple[asyncio.DatagramTransport, SimulatorProtocol]:
    """Start the UDP simulator and return (transport, protocol).

    Args:
        host: Bind address. Use "0.0.0.0" for all interfaces or "127.0.0.1" for loopback.
        port: UDP port to listen on. Defaults to WLED_PORT (21324).
              Pass 0 to let the OS assign a free port (useful in tests).

    Returns:
        Tuple of (transport, protocol). Call transport.close() to stop.
    """
    loop = asyncio.get_running_loop()
    transport, protocol = await loop.create_datagram_endpoint(
        SimulatorProtocol,
        local_addr=(host, port),
    )
    actual_port = transport.get_extra_info("sockname")[1]
    print(f"[SIM] Listening on {host}:{actual_port}")
    return transport, protocol


if __name__ == "__main__":
    async def _main() -> None:
        transport, protocol = await run_simulator()
        try:
            await asyncio.sleep(float("inf"))
        except KeyboardInterrupt:
            pass
        finally:
            transport.close()

    asyncio.run(_main())