import jack # type: ignore
import time
import numpy as np # type: ignore
from pydub import AudioSegment # type: ignore
import random
import sys
import json
import os
import paho.mqtt.client as mqtt # type: ignore
from paho.mqtt.client import Client # type: ignore
import re

# === load and shuffle .m3u playlist ===
def load_playlist(m3u_path):
    with open(m3u_path, "r") as f:
        lines = [line.strip() for line in f if line.strip() and not line.startswith("#")]
    random.shuffle(lines)
    return lines

# === load mp3 song file ===
def decode_mp3(path):
    song = AudioSegment.from_file(path)
    song = song.set_frame_rate(48000).set_channels(2).set_sample_width(4)
    samples = np.array(song.get_array_of_samples()).astype(np.float32) / (2**31)
    return samples.reshape((-1,2)) #stereo:L R

def update_now_playing(playlist_name, song_title):
    path = "/home/media/Music/now_playing.json"
    try:
        with open(path, "r") as f:
            data = json.load(f)
    except FileNotFoundError:
        data = {}

    data[playlist_name.lower()] = os.path.basename(song_title)
    with open(path, "w") as f:
        json.dump(data, f, indent=2)

    # MQTT title broadcast
    mqtt_topic = f"music/now_playing/{playlist_name.lower()}"
    mqtt_payload = os.path.basename(song_title)
    mqtt_payload = os.path.splitext(mqtt_payload)[0]
    mqtt_payload = re.sub(r'\s*\[[^\]]+\]$', '', mqtt_payload)

    try:
        mqtt_client = mqtt.Client(protocol=mqtt.MQTTv311)
        mqtt_client.connect("10.40.0.2", 1883)
        result = mqtt_client.publish(mqtt_topic, mqtt_payload, retain=True)
        result.wait_for_publish()
        if result.rc != mqtt.MQTT_ERR_SUCCESS:
            raise Exception(f"Publish failed with code {result.rc}")
        mqtt_client.disconnect()
    except Exception as e:
        print(f"[MQTT] failed to publish: {e}")

# === Get playlist name from CLI ===
if len(sys.argv) < 2:
    print ("Usage: python3 jack_mp3_player.py Praise|Worship|Instrumental")
    sys.exit(1)

playlist_name = sys.argv[1]
playlist_file = f"/home/media/Music/{playlist_name}.m3u"

# === jack setup ===
client = jack.Client(playlist_name)
out_l = client.outports.register("left")
out_r = client.outports.register("right")

playlist = load_playlist(playlist_file)
current_index = 0
samples = decode_mp3(playlist[current_index])
buffer_pos = [0]

@client.set_process_callback
def process(frames):
    bufsize = frames
    end = buffer_pos[0] + bufsize

    #zero fills if out of data
    if buffer_pos[0] >= len(samples):
        out_l.get_buffer()[:] = np.zeros(bufsize, dtype=np.float32)
        out_r.get_buffer()[:] = np.zeros(bufsize, dtype=np.float32)
        return
    
    chunk = samples[buffer_pos[0]:end]
    #pad if last chunk
    if len(chunk) < bufsize:
        pad = np.zeros((bufsize - len(chunk), 2), dtype=np.float32)
        chunk = np.vstack((chunk, pad))

    out_l.get_buffer()[:] = chunk[:, 0].copy()
    out_r.get_buffer()[:] = chunk[:, 1].copy()
    buffer_pos[0] += bufsize

# === Activate jack ===
with client:
    print(f"Playing playlist: {playlist_name}")
    print(f"Playing song: {playlist[current_index]}")
    update_now_playing(playlist_name, playlist[current_index])
    while True:
        if buffer_pos[0] >= len(samples):
            current_index += 1
            if current_index >= len(playlist):
                playlist = load_playlist(playlist_file)
                current_index = 0

            print(f"Playing: {playlist[current_index]}")
            update_now_playing(playlist_name, playlist[current_index])
            samples = decode_mp3(playlist[current_index])
            buffer_pos[0] = 0
        
        time.sleep(0.1)