#!/usr/bin/env python3
"""Listen for MIDI events via UDP and display note activity on LED matrix.

Non-blocking I/O with immediate MIDI processing and fast frame rendering.
"""
import socket, msgpack, time, os, select, collections

VIZ_PORT = int(os.environ.get("VIZ_PORT", "8082"))
ROUTER_SOCK = os.environ.get("ROUTER_SOCK", "/var/run/arduino-router.sock")
DECAY_FRAMES = int(os.environ.get("DECAY_FRAMES", "8"))
FPS = 30

COLS = 13
ROWS = 8


def pack_frame(pixels):
    bits = 0
    for row in range(ROWS):
        for col in range(COLS):
            if pixels[row][col]:
                bits |= 1 << (row * COLS + col)
    return (bits & 0xFFFFFFFF, (bits >> 32) & 0xFFFFFFFF,
            (bits >> 64) & 0xFFFFFFFF, (bits >> 96) & 0xFFFFFFFF)


def connect_router():
    router = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    router.setblocking(False)
    try:
        router.connect(ROUTER_SOCK)
    except OSError:
        pass
    return router


def main():
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind(("0.0.0.0", VIZ_PORT))
    sock.setblocking(False)

    router = connect_router()
    packer = msgpack.Packer()

    active_notes = {}
    last_frame = None
    buf = b""
    msg_id = 0
    frame_interval = 1.0 / FPS
    next_frame = time.monotonic()
    router_buf = b""
    frames_sent = 0
    frames_dropped = 0

    poller = select.poll()
    poller.register(sock, select.POLLIN)

    while True:
        # Always poll UDP with short timeout
        timeout = max(0, int((next_frame - time.monotonic()) * 1000))
        timeout = min(timeout, 10)  # cap at 10ms for responsiveness

        try:
            events = poller.poll(timeout)
        except OSError:
            events = []

        # Process all pending UDP data immediately
        for fd, flags in events:
            if fd == sock.fileno() and flags & select.POLLIN:
                try:
                    while True:
                        data, _ = sock.recvfrom(1024)
                        buf += data
                except BlockingIOError:
                    pass

                # Parse all complete messages
                while b";" in buf:
                    msg, buf = buf.split(b";", 1)
                    msg = msg.decode().strip()
                    parts = msg.split()
                    if len(parts) >= 3 and parts[0] == "midi":
                        status = int(parts[1])
                        note = int(parts[2])
                        vel = int(parts[3]) if len(parts) >= 4 else 0
                        msg_type = status & 0xF0
                        if msg_type == 0x90 and vel > 0:
                            active_notes[note] = DECAY_FRAMES
                        elif msg_type == 0x80 or (msg_type == 0x90 and vel == 0):
                            active_notes.pop(note, None)

        # Also drain any pending router responses
        if router.fileno() >= 0:
            try:
                while True:
                    router_buf += router.recv(4096)
            except BlockingIOError:
                pass

        # Render frame when it's time
        now = time.monotonic()
        if now >= next_frame:
            next_frame = now + frame_interval

            pixels = [[0] * COLS for _ in range(ROWS)]
            expired = []
            for note, decay in active_notes.items():
                if decay <= 0:
                    expired.append(note)
                    continue
                active_notes[note] = decay - 1
                note_min, note_max = 36, 96
                row = 7 - int((note - note_min) / (note_max - note_min) * 7)
                row = max(0, min(7, row))
                col = note % 12
                pixels[row][col] = 1

            for note in expired:
                del active_notes[note]

            frame = pack_frame(pixels)
            if frame != last_frame:
                msg_id = (msg_id + 1) % 65535
                params = [str(v) for v in frame]
                req = packer.pack([0, msg_id, "draw_frame", params])
                try:
                    router.sendall(req)
                    frames_sent += 1
                except (BlockingIOError, OSError):
                    frames_dropped += 1
                    # Reconnect on error
                    try:
                        router.close()
                    except OSError:
                        pass
                    router = connect_router()
                last_frame = frame


if __name__ == "__main__":
    main()