led2/firmware/main.py

"""LightSync firmware — main entry point.

Two asyncio tasks:
  1. udp_listener_task: non-blocking UDP socket, parses packets, updates state
  2. render_task: 60fps loop, renders animations or writes frame data to strip

Modes:
  - 'idle': strip is off, waiting for first packet
  - 'animation': running a locally-rendered animation (from 0xAC command)
  - 'frame': displaying raw pixel data (from DRGB/DRGBW/DNRGB packets)

Frame timeout: if no frame packet arrives for 2 seconds, reverts to last animation.
"""
import asyncio
import socket
import select
import time
import json

from protocol import (
    classify_packet, parse_animation_cmd,
    parse_drgb, parse_drgbw, parse_dnrgb, FrameAssembler
)
from led_driver import create_strip
from animations import create_animation

# --- Global state ---
config = None
strip = None
current_mode = 'idle'           # 'idle', 'animation', 'frame'
current_animation = None        # dict from parse_animation_cmd
current_anim_instance = None    # animation object with render()
animation_start_ms = 0
frame_pixels = None             # list of (R,G,B) or (R,G,B,W) tuples
frame_is_rgbw = False           # True if last frame was DRGBW
last_frame_ms = 0               # ticks_ms of last frame packet
frame_assembler = None          # FrameAssembler for DNRGB
FRAME_TIMEOUT_MS = 2000         # revert to animation after 2s without frames


def load_config():
    try:
        with open('config.json') as f:
            return json.load(f)
    except OSError:
        return {'led_count': 300, 'strip_type': 'sk6812', 'led_pin': 5,
                'spi_id': 1, 'udp_port': 21324}


async def udp_listener_task(port):
    """Non-blocking UDP listener using select.poll(0) inside asyncio."""
    global current_mode, current_animation, current_anim_instance
    global animation_start_ms, frame_pixels, frame_is_rgbw
    global last_frame_ms, frame_assembler

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.bind(('0.0.0.0', port))
    sock.setblocking(False)
    poller = select.poll()
    poller.register(sock, select.POLLIN)

    print('[udp] Listening on port', port)

    while True:
        events = poller.poll(0)
        if events:
            try:
                data, addr = sock.recvfrom(1500)
                ptype = classify_packet(data)

                if ptype == 'anim_cmd':
                    cmd = parse_animation_cmd(data)
                    if cmd:
                        current_animation = cmd
                        current_anim_instance = create_animation(cmd)
                        current_mode = 'animation'
                        animation_start_ms = time.ticks_ms()
                        print('[udp] Animation:', cmd['animation'])

                elif ptype == 'drgb':
                    pixels = parse_drgb(data)
                    frame_pixels = pixels
                    frame_is_rgbw = False
                    current_mode = 'frame'
                    last_frame_ms = time.ticks_ms()

                elif ptype == 'drgbw':
                    pixels = parse_drgbw(data)
                    frame_pixels = pixels
                    frame_is_rgbw = True
                    current_mode = 'frame'
                    last_frame_ms = time.ticks_ms()

                elif ptype == 'dnrgb':
                    start_idx, pixels = parse_dnrgb(data)
                    now = time.ticks_ms()
                    complete = frame_assembler.feed(start_idx, pixels, now)
                    if complete:
                        frame_pixels = complete
                        frame_is_rgbw = False  # DNRGB is always RGB
                        current_mode = 'frame'
                        last_frame_ms = now

                else:
                    print('[udp] Unknown packet type:', data[0] if data else '?')
            except OSError:
                pass

        await asyncio.sleep_ms(5)


async def render_task():
    """Animation/frame render loop at ~60fps."""
    global current_mode, frame_pixels

    while True:
        if current_mode == 'frame' and frame_pixels:
            # Check frame timeout — revert to animation if no packets for 2s
            if time.ticks_diff(time.ticks_ms(), last_frame_ms) > FRAME_TIMEOUT_MS:
                if current_anim_instance:
                    current_mode = 'animation'
                    animation_start_ms = time.ticks_ms()
                    print('[render] Frame timeout, reverting to animation')
                else:
                    current_mode = 'idle'
                    strip.clear()
            else:
                # Write frame to strip
                if frame_is_rgbw:
                    strip.write_rgbw(frame_pixels)
                else:
                    strip.write_rgb(frame_pixels)

        elif current_mode == 'animation' and current_anim_instance:
            t = time.ticks_diff(time.ticks_ms(), animation_start_ms) / 1000.0
            pixels = current_anim_instance.render(t, config['led_count'])
            strip.write_rgb(pixels)

        await asyncio.sleep_ms(16)  # ~60 fps


async def main():
    global config, strip, frame_assembler

    config = load_config()
    strip = create_strip(config)
    strip.clear()
    frame_assembler = FrameAssembler(config['led_count'])

    print('[main] LightSync firmware ready')
    print('[main] Strip:', config['strip_type'], '/', config['led_count'], 'LEDs')
    print('[main] UDP port:', config['udp_port'])

    asyncio.create_task(udp_listener_task(config['udp_port']))
    asyncio.create_task(render_task())

    while True:
        await asyncio.sleep(1)

asyncio.run(main())