led-sync-studio/src/led_sync/transport/udp_client.py

"""
UDP transport client for ESP32 LED controller.
Protocol: docs/protocol.md — v1, port 4210, fire-and-forget JSON datagrams.

Decisions applied:
  D-10: asyncio DatagramProtocol, fire-and-forget sendto()
  D-11: Optional periodic heartbeat (status ping), non-blocking
  D-12: connected/disconnected tracked via STATUS responses
  D-14: Minimal error handling (fire-and-forget transport)
"""
import asyncio
import json
import logging
from typing import Callable

logger = logging.getLogger(__name__)

PROTOCOL_VERSION = 1
MAX_PAYLOAD_BYTES = 512


class ESP32Transport(asyncio.DatagramProtocol):
    """
    asyncio DatagramProtocol client for the ESP32 LED controller.

    Usage:
        transport = await create_esp32_transport("192.168.1.100")
        transport.send_command({"zone": "wand", "animation": "chase", "params": {"speed": 0.5}})

    Public interface:
        send_command(cmd: dict) -> None  — fire-and-forget, non-blocking
        connected: bool                  — True after first STATUS ok response
        last_status: dict | None         — most recent STATUS response from ESP32
    """

    def __init__(self, on_status: Callable[[dict], None] | None = None):
        self._transport: asyncio.DatagramTransport | None = None
        self.connected: bool = False
        self.last_status: dict | None = None
        self._on_status = on_status  # optional callback invoked on STATUS ok

    # -------------------------------------------------------------------------
    # asyncio DatagramProtocol callbacks
    # -------------------------------------------------------------------------

    def connection_made(self, transport: asyncio.DatagramTransport) -> None:
        """Called by asyncio when the UDP socket is bound to the remote address."""
        self._transport = transport
        logger.debug("UDP socket bound to ESP32 remote address")

    def datagram_received(self, data: bytes, addr: tuple) -> None:
        """
        Called when a UDP datagram arrives from the ESP32.

        Parses STATUS ok responses to update connection state (D-12).
        All other messages and malformed JSON are silently ignored (D-14).
        """
        try:
            msg = json.loads(data.decode("utf-8"))
        except Exception:
            logger.debug("Received non-JSON datagram from %s — ignored", addr)
            return

        if msg.get("status") == "ok":
            self.connected = True
            self.last_status = msg
            logger.debug("ESP32 STATUS ok: %s", msg)
            if self._on_status is not None:
                self._on_status(msg)

    def error_received(self, exc: Exception) -> None:
        """
        Called on non-fatal UDP errors (e.g., ICMP port unreachable).
        Fire-and-forget transport — errors are logged but not re-raised.
        """
        logger.debug("UDP error (non-fatal): %s", exc)

    def connection_lost(self, exc: Exception | None) -> None:
        """Called when the UDP socket is closed."""
        self.connected = False
        self._transport = None
        logger.debug("UDP transport closed (exc=%s)", exc)

    # -------------------------------------------------------------------------
    # Public API
    # -------------------------------------------------------------------------

    def send_command(self, cmd: dict) -> None:
        """
        Fire-and-forget: inject "v":1 and send cmd as a JSON UDP datagram.

        Safe to call from the asyncio event loop thread at any time.
        Silent drop (no exception) if:
          - Transport not yet ready (pre-connection)
          - Resulting payload exceeds MAX_PAYLOAD_BYTES (512)

        Matches docs/protocol.md exactly — "v":1 is injected here so callers
        do not need to include it.
        """
        if self._transport is None:
            logger.debug("send_command() called before transport ready — dropped")
            return

        payload = {"v": PROTOCOL_VERSION, **cmd}
        data = json.dumps(payload).encode("utf-8")

        if len(data) > MAX_PAYLOAD_BYTES:
            logger.warning(
                "Command payload %d bytes exceeds %d-byte limit — dropped",
                len(data),
                MAX_PAYLOAD_BYTES,
            )
            return

        self._transport.sendto(data)

    def send_status_ping(self) -> None:
        """
        Send a STATUS query to the ESP32 (D-11 heartbeat).

        Fire-and-forget — a STATUS ok response via datagram_received() will
        update connected=True and last_status. Silent drop if not ready.
        """
        self.send_command({"cmd": "status"})


async def create_esp32_transport(
    host: str,
    port: int = 4210,
    on_status: Callable[[dict], None] | None = None,
) -> "ESP32Transport":
    """
    Async factory: create and return an ESP32Transport bound to host:port.

    Sends an initial STATUS ping to establish connection state (D-12).
    The reply (if ESP32 is reachable) will set connected=True via
    datagram_received().

    Args:
        host: ESP32 IP address (e.g., "192.168.1.100")
        port: UDP port (default 4210 per docs/protocol.md)
        on_status: optional callback invoked on every STATUS ok response

    Returns:
        ESP32Transport instance ready for send_command() calls
    """
    loop = asyncio.get_running_loop()
    _, protocol = await loop.create_datagram_endpoint(
        lambda: ESP32Transport(on_status=on_status),
        remote_addr=(host, port),
    )
    # Initial status ping — non-blocking, reply updates connected state (D-12)
    protocol.send_status_ping()
    return protocol