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

Non-blocking I/O with proper POLLOUT queue, instrumentation, and fast 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", "3"))
MAX_NOTE_AGE = int(os.environ.get("MAX_NOTE_AGE", "15"))

COLS = 13
ROWS = 8

frame_interval = 1.0 / 30.0  # 30 FPS target
next_frame_time = 0.0

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

SOCK_BUFFER = 4096
ROUTER_RECONNECT_DELAY = 0.1


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)


class FrameQueue:
    def __init__(self, maxsize=10):
        self._q = collections.deque()
        self._maxsize = maxsize

    def push(self, frame_bytes):
        if len(self._q) >= self._maxsize:
            return False  # dropped
        self._q.append(frame_bytes)
        return True

    def pop(self):
        return self._q.popleft() if self._q else None

    def __len__(self):
        return len(self._q)

    def full(self):
        return len(self._q) >= self._maxsize


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():
    global next_frame_time

    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 = {}
    last_frame = None
    buf = b""
    msg_id = 0
    next_frame_time = time.monotonic()

    pending_frames = FrameQueue(maxsize=5)
    router_fd = router.fileno()
    _poll.register(router, POLLIN | POLLOUT)

    frames_sent = 0
    frames_dropped = 0
    router_errors = 0
    stats_last_write = time.monotonic()
    fps_sample_start = time.monotonic()
    fps_sample_frames = 0

    while True:
        now = time.monotonic()

        timeout = max(0, int((next_frame_time - now) * 1000))
        timeout = min(timeout, 5)  # 5ms max to stay responsive

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

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

                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": DECAY_FRAMES, "age": 0}
                        elif msg_type == 0x80 or (msg_type == 0x90 and vel == 0):
                            active_notes.pop(note, None)

            # Process router readable
            if fd == router_fd and (flags & POLLIN):
                try:
                    while True:
                        d = router.recv(SOCK_BUFFER)
                        if not d:
                            break
                except BlockingIOError:
                    pass

            # Process router writable — drain pending frames
            if fd == router_fd and (flags & POLLOUT):
                frame = pending_frames.pop()
                if frame:
                    try:
                        router.sendall(frame)
                        frames_sent += 1
                    except (BlockingIOError, OSError):
                        frames_dropped += 1
                        router_errors += 1
                        pending_frames.push(frame)
                        time.sleep(ROUTER_RECONNECT_DELAY)
                        old_fd = router_fd
                        try:
                            router.close()
                        except OSError:
                            pass
                        router = connect_router()
                        router_fd = router.fileno()
                        try:
                            _poll.unregister(old_fd)
                        except (KeyError, OSError):
                            pass
                        _poll.register(router, POLLIN | POLLOUT)

        # Render frame when due
        if now >= next_frame_time:
            next_frame_time = now + frame_interval
            fps_sample_frames += 1

            pixels = [[0] * COLS for _ in range(ROWS)]
            expired = []
            for note, state in list(active_notes.items()):
                state["age"] += 1
                state["decay"] -= 1
                if state["decay"] <= 0 or state["age"] > MAX_NOTE_AGE:
                    expired.append(note)
                    continue
                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])
                if pending_frames.push(req):
                    pass  # queued
                else:
                    frames_dropped += 1
                last_frame = frame

            # Write stats periodically (every 2s)
            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, pending_frames, fps, len(active_notes))
                stats_last_write = now
                fps_sample_start = now
                fps_sample_frames = 0


def _write_stats(sent, dropped, pending, fps, active_count):
    import json
    try:
        with open(STATS_FILE, "w") as f:
            json.dump({
                "frames_sent": sent,
                "frames_dropped": dropped,
                "pending_queue": len(pending),
                "queue_full": pending.full(),
                "fps": round(fps, 1),
                "active_notes": active_count,
                "ts": time.time()
            }, f)
    except OSError:
        pass


if __name__ == "__main__":
    main()