import time
import numpy as np
import sounddevice as sd
import pyloudnorm as pyln
import paho.mqtt.client as mqtt

# MQTT
MQTT_HOST = "10.40.0.2"
MQTT_PORT = 1883

TOPIC_M = "stream/audio/lufs/momentary"
TOPIC_S = "stream/audio/lufs/shortterm"
TOPIC_I = "stream/audio/lufs/integrated"
DEVICE = "default"

# Audio
RATE = 48000
CHANNELS = 2

def find_scarlett():
    for i, d in enumerate(sd.query_devices()):
        name = d["name"].lower()
        if "scarlett" in name and d["max_input_channels"] >= 2:
            return i
    raise RuntimeError("Scarlett stereo input not found")

# Windows
M_WINDOW = 0.4
S_WINDOW = 3.0
UPDATE = 0.25

meter = pyln.Meter(RATE)

integrated_buffer = np.array([], dtype=np.float32)
shortterm_buffer = np.array([], dtype=np.float32)

RESET_TOPIC = "stream/audio/lufs/reset"

reset_requested = False


def on_message(client, userdata, msg):
    global reset_requested

    if msg.topic == RESET_TOPIC:
        print("Resetting integrated LUFS")
        reset_requested = True

client = mqtt.Client()

client.on_message = on_message

client.connect(MQTT_HOST, MQTT_PORT)

client.subscribe(RESET_TOPIC)

client.loop_start()

print("LUFS → MQTT running")

while True:

    try:
        samples = sd.rec(
            int(UPDATE * RATE),
            samplerate=RATE,
            channels=CHANNELS,
            device=DEVICE,
            dtype="float32",
        )
        sd.wait()

    except Exception as e:
        print(f"Audio route changed / device unavailable: {e}")
        time.sleep(2)
        continue

    mono = samples.flatten()

    if reset_requested:
        integrated_buffer = np.array([], dtype=np.float32)
        reset_requested = False

    integrated_buffer = np.concatenate(
        [integrated_buffer, mono]
    )[-RATE * 60 * 120:]

    shortterm_buffer = np.concatenate(
        [shortterm_buffer, mono]
    )[-int(S_WINDOW * RATE):]

    try:

        m_samples = shortterm_buffer[
            -int(M_WINDOW * RATE):
        ]

        momentary = meter.integrated_loudness(
            m_samples
        )

        shortterm = meter.integrated_loudness(
            shortterm_buffer
        )

        integrated = meter.integrated_loudness(
            integrated_buffer
        )

        client.publish(
            TOPIC_M,
            f"{momentary:.1f}",
            retain=True
        )

        client.publish(
            TOPIC_S,
            f"{shortterm:.1f}",
            retain=True
        )

        client.publish(
            TOPIC_I,
            f"{integrated:.1f}",
            retain=True
        )

        print(
            f"\rM {momentary:.1f} "
            f"S {shortterm:.1f} "
            f"I {integrated:.1f}",
            end=""
        )

    except Exception:
        pass