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

Uses select() for polling both UDP and router RPC responses.
"""
import socket, msgpack, time, os, select

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 = 15

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 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 = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    router.setblocking(False)
    router.connect(ROUTER_SOCK)
    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""

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

    while True:
        timeout = max(0, int((next_frame - time.monotonic()) * 1000))
        try:
            events = poller.poll(timeout)
        except OSError:
            events = []

        for fd, flags in events:
            if fd == sock.fileno() and flags & select.POLLIN:
                try:
                    data, _ = sock.recvfrom(1024)
                    buf += data
                    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)
                except BlockingIOError:
                    pass

            if fd == router.fileno() and flags & select.POLLIN:
                try:
                    router_buf += router.recv(4096)
                except BlockingIOError:
                    pass

        if time.monotonic() >= next_frame:
            next_frame = time.monotonic() + frame_interval

            pixels = [[0] * COLS for _ in range(ROWS)]
            for note, decay in list(active_notes.items()):
                if decay <= 0:
                    del active_notes[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

            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)
                except BlockingIOError:
                    pass
                last_frame = frame

if __name__ == "__main__":
    main()