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

Event-driven: renders and sends frame immediately on each MIDI event.
Background 5 FPS refresh catches any lost frames from the router.
"""
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")

COLS = 13
ROWS = 8

REFRESH_INTERVAL = 0.2

STATS_FILE = "/tmp/viz-stats.json"
_poll = select.poll()
POLLIN = select.POLLIN

SOCK_BUFFER = 4096


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)
    w0 = bits & 0xFFFFFFFF
    w1 = (bits >> 32) & 0xFFFFFFFF
    w2 = (bits >> 64) & 0xFFFFFFFF
    w3 = (bits >> 96) & 0xFFFFFFFF
    return [str(w0), str(w1), str(w2), str(w3)]


def render(active_notes):
    pixels = [[0] * COLS for _ in range(ROWS)]
    note_min, note_max = 36, 96
    for note in active_notes:
        row = 7 - int((note - note_min) / (note_max - note_min) * 7)
        row = max(0, min(7, row))
        col = note % 12
        pixels[row][col] = 1
    return pack_frame(pixels)


def connect_router():
    router = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    router.settimeout(0.3)
    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()

    _poll.register(sock, POLLIN)

    active_notes = set()
    last_sent_params = None
    pending_bytes = None
    buf = b""
    msg_id = 0

    frames_sent = 0
    frames_dropped = 0
    bytes_received = 0

    last_refresh = time.monotonic()
    stats_last_write = time.monotonic()
    fps_sample_start = time.monotonic()
    fps_sample_frames = 0
    events_processed = 0
    raw_packets = 0

    while True:
        now = time.monotonic()
        refresh_due = (now - last_refresh) >= REFRESH_INTERVAL

        try:
            events = _poll.poll(5)
        except OSError:
            events = []

        for fd, flags in events:
            if fd == sock.fileno() and (flags & POLLIN):
                try:
                    while True:
                        data, addr = sock.recvfrom(SOCK_BUFFER)
                        buf += data
                        bytes_received += len(data)
                        raw_packets += 1
                except BlockingIOError:
                    pass

                while b";" in buf:
                    raw_msg, buf = buf.split(b";", 1)
                    msg = raw_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
                        changed = False
                        if msg_type == 0x90 and vel > 0:
                            if note not in active_notes:
                                active_notes.add(note)
                                changed = True
                        elif msg_type == 0x80 or (msg_type == 0x90 and vel == 0):
                            if note in active_notes:
                                active_notes.discard(note)
                                changed = True
                        if changed:
                            events_processed += 1
                            params = render(active_notes)
                            if params != last_sent_params:
                                msg_id = (msg_id + 1) % 65535
                                pending_bytes = packer.pack([0, msg_id, "draw_frame", params])
                                last_sent_params = params

        # Try to send pending frame — no POLLOUT dependency
        if pending_bytes is not None:
            try:
                router.sendall(pending_bytes)
                pending_bytes = None
                frames_sent += 1
            except (BlockingIOError, OSError):
                frames_dropped += 1

        # Background refresh: resend last params to recover from router drops
        if refresh_due:
            last_refresh = now
            fps_sample_frames += 1
            if last_sent_params is not None and pending_bytes is None:
                msg_id = (msg_id + 1) % 65535
                refresh_bytes = packer.pack([0, msg_id, "draw_frame", last_sent_params])
                try:
                    router.sendall(refresh_bytes)
                    frames_sent += 1
                except (BlockingIOError, OSError):
                    pass

            if now - stats_last_write >= 2.0:
                fps = fps_sample_frames / (now - fps_sample_start) if (now - fps_sample_start) > 0 else 0
                _write_stats(frames_sent, frames_dropped, fps, len(active_notes), events_processed, bytes_received, raw_packets)
                stats_last_write = now
                fps_sample_start = now
                fps_sample_frames = 0
                events_processed = 0


def _write_stats(sent, dropped, fps, active_count, events, bytes_rcv, packets):
    import json
    try:
        with open(STATS_FILE, "w") as f:
            json.dump({
                "frames_sent": sent,
                "frames_dropped": dropped,
                "fps": round(fps, 1),
                "active_notes": active_count,
                "midi_events": events,
                "bytes_received": bytes_rcv,
                "raw_packets": packets,
                "ts": time.time()
            }, f)
    except OSError:
        pass


if __name__ == "__main__":
    main()