import jack
import socket
import numpy as np
import time

PORT = 5004
JACK_CLIENT_NAME = "Player"
BUFFER_SIZE = 4096

sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
sock.bind(("0.0.0.0", PORT))
sock.setblocking(False)

client = jack.Client(JACK_CLIENT_NAME)

out_l = client.outports.register("out_l")
out_r = client.outports.register("out_r")

buffer = np.zeros((2, 0), dtype=np.float32)

def fetch_samples():
    global buffer
    try:
        data, _ = sock.recvfrom(BUFFER_SIZE)
        int16 = np.frombuffer(data, dtype=np.int16)
        interleaved = int16.reshape(-1, 2).T
        float_audio = interleaved.astype(np.float32) / 32767.0
        buffer = np.concatenate((buffer, float_audio), axis=1)
    except BlockingIOError:
        pass

@client.set_process_callback
def process(frames):
    global buffer
    fetch_samples()
    needed = frames

    if buffer.shape[1] >= needed:
        chunk = buffer[:, :needed]
        buffer = buffer[:, needed:]
    else:
        chunk = np.zeros((2, needed), dtype=np.float32)
        if buffer.shape[1] > 0:
            chunk[:, :buffer.shape[1]] = buffer
            buffer = np.zeros((2, 0), dtype=np.float32)

    out_l.get_array()[:] = chunk[0]
    out_r.get_array()[:] = chunk[1]

with client:
    print(f"Receving on UDP port {PORT} and playing via JACK")
    try:
        while True:
            time.sleep(0.1)
    except KeyboardInterrupt:
        print("Exiting...")