led-controller/src/util/espnow_registry.py

"""Handle ESP-NOW uplink from bridge and push group membership."""

from __future__ import annotations

import json
from typing import Any, Dict, Optional

from models.device import Device, normalize_mac  # noqa: F401 — re-export for callers
from models.group import Group
from models.transport import get_current_bridge
from util.bridge_envelope import build_groups_envelope
from util.espnow_ping import record_ping_rsp
from util.espnow_wire import (
    MSG_ANNOUNCE,
    MSG_PING_RSP,
    WIRE_MAGIC,
    mac_bytes_to_hex,
    parse_announce,
    wire_msg_type,
)
from util.groups_for_device import groups_for_mac


async def handle_bridge_uplink(peer_mac: bytes, payload: bytes) -> None:
    """Dispatch binary wire or JSON v1 hello from bridge uplink."""
    if not payload:
        return
    if payload[0] == WIRE_MAGIC:
        mt = wire_msg_type(payload)
        if mt == MSG_ANNOUNCE:
            await handle_espnow_announce(peer_mac, payload)
        elif mt == MSG_PING_RSP:
            record_ping_rsp(peer_mac, payload)
        return
    if payload[:1] == b"{":
        try:
            data = json.loads(payload.decode("utf-8"))
        except (UnicodeError, ValueError, TypeError):
            return
        if isinstance(data, dict):
            await handle_json_hello(peer_mac, data)


async def _after_device_registered(mac_hex: str) -> None:
    await push_groups_to_mac(mac_hex)


async def handle_json_hello(peer_mac: bytes, data: Dict[str, Any]) -> None:
    """Register device from driver JSON boot hello."""
    if data.get("v") != "1":
        return
    mac_hex = mac_bytes_to_hex(peer_mac)
    if not mac_hex:
        return

    name = data.get("name")
    nested = data.get("settings")
    if not name and isinstance(nested, dict):
        name = nested.get("name")
    name = str(name or mac_hex).strip() or mac_hex

    num_leds = None
    color_order = None
    startup_mode = None
    brightness = None
    if isinstance(nested, dict):
        try:
            num_leds = int(nested.get("num_leds")) if nested.get("num_leds") is not None else None
        except (TypeError, ValueError):
            pass
        color_order = nested.get("color_order")
        startup_mode = nested.get("startup_mode")
        try:
            brightness = int(nested.get("brightness")) if nested.get("brightness") is not None else None
        except (TypeError, ValueError):
            pass

    devices = Device()
    did, persisted = devices.upsert_espnow_announced(
        mac_hex,
        name,
        device_type=data.get("type", "led"),
        num_leds=num_leds,
        color_order=color_order,
        startup_mode=startup_mode,
        brightness=brightness,
    )
    if not did:
        return
    if persisted:
        print(f"[espnow] registered mac={did} name={name!r} (json hello)")
    await _after_device_registered(mac_hex)


async def handle_espnow_announce(peer_mac: bytes, packet: bytes) -> None:
    info = parse_announce(packet)
    if not info:
        return
    mac_hex = mac_bytes_to_hex(peer_mac)
    if not mac_hex:
        return

    devices = Device()
    did, persisted = devices.upsert_espnow_announced(
        mac_hex,
        info["name"],
        device_type=info.get("device_type", "led"),
        num_leds=info.get("num_leds"),
        color_order=info.get("color_order"),
        startup_mode=info.get("startup_mode"),
        brightness=info.get("brightness"),
    )
    if not did:
        return
    if persisted:
        print(f"[espnow] registered mac={did} name={info['name']!r}")
    await _after_device_registered(mac_hex)


async def push_groups_for_group_devices(gdoc: dict) -> None:
    """Push group membership to each device MAC listed on the group."""
    if not isinstance(gdoc, dict):
        return
    mac_list = gdoc.get("devices") if isinstance(gdoc.get("devices"), list) else []
    for mac in mac_list:
        m = normalize_mac(str(mac))
        if m:
            await push_groups_to_mac(m)


async def push_groups_broadcast() -> bool:
    """No aggregate broadcast for group assignment; use per-device push."""
    return False


async def push_groups_all_espnow_devices() -> Dict[str, Any]:
    """Push ``set_groups`` envelopes to every ESP-NOW device in the registry."""
    devices_model = Device()
    macs: list[str] = []
    skipped = 0
    for did, doc in devices_model.items():
        if str(doc.get("transport") or "espnow").strip().lower() != "espnow":
            continue
        mac = normalize_mac(str(did)) or normalize_mac(str(doc.get("address") or ""))
        if not mac:
            skipped += 1
            continue
        macs.append(mac)
    sent = 0
    failed = 0
    for mac in macs:
        if await push_groups_to_mac(mac):
            sent += 1
        else:
            failed += 1
    ok = bool(macs) and failed == 0
    return {
        "ok": ok,
        "sent": sent,
        "failed": failed,
        "skipped": skipped,
        "total": len(macs),
    }


async def push_groups_to_mac(mac_hex: str) -> bool:
    """Unicast groups envelope to one driver (set_groups true)."""
    mac = normalize_mac(mac_hex)
    if not mac:
        return False
    gids = groups_for_mac(mac, Group())
    bridge = get_current_bridge()
    if bridge is None:
        return False
    envelope = build_groups_envelope(mac, gids)
    ok = await bridge.send(envelope)
    if ok:
        print(f"[espnow] groups sent mac={mac} groups={gids!r}")
    return bool(ok)