# audio_engine.py

import numpy as np
import soundfile as sf
import jack
import threading
import subprocess
import os
import pyloudnorm as pyln


LUFS_TARGET = -14.0
SOUNDFILE_FORMATS = {'.wav', '.flac', '.ogg', '.aiff', '.aif'}
FFMPEG_FORMATS    = {'.mp3', '.m4a', '.mp4', '.opus', '.aac', '.wma'}

NUM_MIXER_CHANNELS = 8
DEFAULT_CHANNEL_NAMES = [
    "Mic 1", "Mic 2", "Aux 1",
    "Aux 2", "Aux 3", "Aux 4", "Aux 5", "Aux 6",
]


def load_audio(filepath: str, target_sr: int) -> np.ndarray:
    ext = os.path.splitext(filepath)[1].lower()

    if ext in SOUNDFILE_FORMATS:
        data, sr = sf.read(filepath, always_2d=True, dtype='float32')
        data = data.T

    elif ext in FFMPEG_FORMATS:
        cmd = [
            'ffmpeg', '-i', filepath, '-f', 'f32le', '-acodec', 'pcm_f32le',
            '-ar', str(target_sr), '-ac', '2', '-', '-loglevel', 'quiet'
        ]
        result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        if result.returncode != 0:
            raise ValueError(f"FFmpeg failed: {result.stderr.decode()}")
        raw  = np.frombuffer(result.stdout, dtype=np.float32)
        data = raw.reshape(-1, 2).T
        sr   = target_sr

    else:
        try:
            data, sr = sf.read(filepath, always_2d=True, dtype='float32')
            data = data.T
        except Exception as e:
            raise ValueError(f"Unsupported format: {ext} — {e}")

    if sr != target_sr:
        import librosa
        data = np.array([
            librosa.resample(data[ch], orig_sr=sr, target_sr=target_sr)
            for ch in range(data.shape[0])
        ], dtype=np.float32)

    if data.shape[0] == 1:
        data = np.vstack([data, data])

    data = np.clip(data, -1.0, 1.0)
    return data


def normalize_lufs(audio: np.ndarray, sr: int,
                   target: float = LUFS_TARGET) -> np.ndarray:
    meter = pyln.Meter(sr)
    loudness = meter.integrated_loudness(audio.T)
    if np.isinf(loudness):
        return audio
    gain_db = target - loudness
    gain_linear = 10 ** (gain_db / 20.0)
    normalized = audio * gain_linear
    return np.clip(normalized, -1.0, 1.0)


class Track:
    def __init__(self, filepath: str, audio: np.ndarray, sr: int):
        self.filepath    = filepath
        self.filename    = os.path.basename(filepath)
        self.audio       = audio
        self.sr          = sr
        self.num_samples = audio.shape[1]
        self.duration    = self.num_samples / sr


class DeckState:
    def __init__(self):
        self.track    = None
        self.queued   = None
        self.position = 0
        self.playing  = False
        self.lock     = threading.Lock()
        self.manual_elapsed = 0.0
        self.manual_duration = 0.0

    def load(self, track: Track):
        with self.lock:
            self.track    = track
            self.position = 0
            self.playing  = False

    def load_queue(self, track: Track):
        with self.lock:
            self.queued = track

    def play(self):
        with self.lock:
            if self.track:
                self.playing = True

    def pause(self):
        with self.lock:
            self.playing = False

    def stop(self):
        with self.lock:
            self.playing  = False
            self.track    = None
            self.position = 0

    def skip_seconds(self, seconds: int):
        with self.lock:
            if self.track:
                new = self.position + int(seconds * self.track.sr)
                self.position = max(0, min(new, self.track.num_samples - 1))

    def go_to_start(self):
        with self.lock:
            self.position = 0

    def go_to_end(self):
        with self.lock:
            if self.track:
                self.position = self.track.num_samples - 1

    def get_elapsed(self) -> str:
        with self.lock:
            if not self.track:
                return "00:00"
            s = self.position / self.track.sr
        return f"{int(s//60):02d}:{int(s%60):02d}"

    def get_remaining(self) -> str:
        with self.lock:
            if not self.track:
                return "-00:00"
            remaining = max(0, self.track.duration - (self.position / self.track.sr))
        return f"-{int(remaining//60):02d}:{int(remaining%60):02d}"


class MixerChannel:
    def __init__(self, name: str, index: int):
        self.name     = name
        self.index    = index
        self.volume   = 1.0
        self.gain     = 1.0      # pre-fader gain trim
        self.muted    = False
        self.solo     = False
        self.pfl     = False
        self.mic_on   = False
        self.vu_peak  = 0.0
        self._vu_smooth = 0.0
        self.in_port_L = None
        self.in_port_R = None


class MasterBus:
    def __init__(self, sr: int):
        self.vu_peak = 0.0
        self._vu_smooth = 0.0
        self.lufs_current = -70.0
        self.lufs_short = -70.0
        self.lufs_long = -70.0
        self.volume = 1.0
        self.pfl_hp_blend = 0.5   # 0.0 = PFL only, 0.5 = equal, 1.0 = HP only
        self.headphone_volume = 0.7
        self._lufs_accum = 0


class AudioEngine:
    def __init__(self):
        self.client = jack.Client("RadioPanel")
        self.sr     = self.client.samplerate

        # Output ports for digital decks (standalone, not mixed into master)
        self.ports = {
            'deck1': (
                self.client.outports.register("deck1_L"),
                self.client.outports.register("deck1_R"),
            ),
            'deck2': (
                self.client.outports.register("deck2_L"),
                self.client.outports.register("deck2_R"),
            ),
            'cart': (
                self.client.outports.register("cart_L"),
                self.client.outports.register("cart_R"),
            ),
        }

        # LUFS metering input (external processed signal)
        self.lufs_in_L = self.client.inports.register("lufs_in_L")
        self.lufs_in_R = self.client.inports.register("lufs_in_R")

        # Mixer channels: 8 stereo input ports
        self.mixer = {
            'channels': [],
            'master': MasterBus(self.sr),
        }
        for i in range(NUM_MIXER_CHANNELS):
            ch = MixerChannel(DEFAULT_CHANNEL_NAMES[i], i)
            ch.in_port_L = self.client.inports.register(f"mixer_ch_{i+1:02d}_in_L")
            ch.in_port_R = self.client.inports.register(f"mixer_ch_{i+1:02d}_in_R")
            self.mixer['channels'].append(ch)

        # Master output ports
        self.master_out = (
            self.client.outports.register("master_out_L"),
            self.client.outports.register("master_out_R"),
        )

        # Headphone output ports (all channels except mics)
        self.headphone_out = (
            self.client.outports.register("headphone_out_L"),
            self.client.outports.register("headphone_out_R"),
        )

        # Studio output ports (same as main, mutes when any mic is on)
        self.studio_out = (
            self.client.outports.register("studio_out_L"),
            self.client.outports.register("studio_out_R"),
        )

        # PFL output ports
        self.pfl_out = (
            self.client.outports.register("pfl_out_L"),
            self.client.outports.register("pfl_out_R"),
        )

        # Deck states (for digital decks and carts)
        self.decks = {
            'deck1': DeckState(),
            'deck2': DeckState(),
            'deck3': DeckState(),
            'deck4': DeckState(),
            'cart1': DeckState(),
            'cart2': DeckState(),
            'cart3': DeckState(),
            'cart4': DeckState(),
        }

        self.client.set_process_callback(self._process)
        self.client.activate()

    def get_channel(self, index: int) -> MixerChannel:
        return self.mixer['channels'][index]

    def set_channel_volume(self, index: int, volume: float):
        if 0 <= index < NUM_MIXER_CHANNELS:
            self.mixer['channels'][index].volume = max(0.0, min(1.0, volume))

    def set_channel_gain(self, index: int, gain: float):
        if 0 <= index < NUM_MIXER_CHANNELS:
            self.mixer['channels'][index].gain = max(0.0, min(2.0, gain))

    def set_channel_mute(self, index: int, muted: bool):
        if 0 <= index < NUM_MIXER_CHANNELS:
            self.mixer['channels'][index].muted = muted

    def set_channel_solo(self, index: int, solo: bool):
        if 0 <= index < NUM_MIXER_CHANNELS:
            self.mixer['channels'][index].solo = solo

    def set_channel_pfl(self, index: int, pfl: bool):
        if 0 <= index < NUM_MIXER_CHANNELS:
            self.mixer['channels'][index].pfl = pfl

    def set_channel_mic_on(self, index: int, on: bool):
        if 0 <= index < NUM_MIXER_CHANNELS:
            self.mixer['channels'][index].mic_on = on

    def get_channel_mic_on(self, index: int) -> bool:
        if 0 <= index < NUM_MIXER_CHANNELS:
            return self.mixer['channels'][index].mic_on
        return False

    def _fill_port(self, port_L, port_R, deck: DeckState, frames: int):
        buf_L = port_L.get_array()
        buf_R = port_R.get_array()
        buf_L.fill(0.0)
        buf_R.fill(0.0)

        with deck.lock:
            playing = deck.playing
            track   = deck.track
            pos     = deck.position

        if not playing or track is None:
            return

        end   = min(pos + frames, track.num_samples)
        n     = end - pos

        if n <= 0:
            with deck.lock:
                deck.playing = False
            return

        buf_L[:n] = track.audio[0, pos:end]
        buf_R[:n] = track.audio[1, pos:end]

        with deck.lock:
            deck.position += n
            if deck.position >= track.num_samples:
                deck.playing  = False
                deck.position = track.num_samples - 1

    def _process(self, frames: int):
        # ── Digital decks ───────────────────────────────────────────────
        self._fill_port(*self.ports['deck1'], self.decks['deck1'], frames)
        self._fill_port(*self.ports['deck2'], self.decks['deck2'], frames)

        # ── Mix carts ───────────────────────────────────────────────────
        buf_L = self.ports['cart'][0].get_array()
        buf_R = self.ports['cart'][1].get_array()
        buf_L.fill(0.0)
        buf_R.fill(0.0)

        for cart_key in ('cart1', 'cart2', 'cart3', 'cart4'):
            deck = self.decks[cart_key]
            with deck.lock:
                playing = deck.playing
                track   = deck.track
                pos     = deck.position
            if not playing or track is None:
                continue
            end   = min(pos + frames, track.num_samples)
            n     = end - pos
            if n <= 0:
                with deck.lock:
                    deck.playing = False
                continue
            buf_L[:n] += track.audio[0, pos:end]
            buf_R[:n] += track.audio[1, pos:end]
            with deck.lock:
                deck.position += n
                if deck.position >= track.num_samples:
                    deck.playing  = False
                    deck.position = track.num_samples - 1

        # ── Mixer ───────────────────────────────────────────────────────
        dt = frames / self.sr
        vu_decay = np.exp(-dt / 0.15)
        dt = frames / self.sr
        vu_decay = np.exp(-dt / 0.15)
        master_buf_L = self.master_out[0].get_array()
        master_buf_R = self.master_out[1].get_array()
        master_buf_L.fill(0.0)
        master_buf_R.fill(0.0)

        headphone_buf_L = self.headphone_out[0].get_array()
        headphone_buf_R = self.headphone_out[1].get_array()
        headphone_buf_L.fill(0.0)
        headphone_buf_R.fill(0.0)

        studio_buf_L = self.studio_out[0].get_array()
        studio_buf_R = self.studio_out[1].get_array()
        studio_buf_L.fill(0.0)
        studio_buf_R.fill(0.0)

        pfl_buf_L = self.pfl_out[0].get_array()
        pfl_buf_R = self.pfl_out[1].get_array()
        pfl_buf_L.fill(0.0)
        pfl_buf_R.fill(0.0)

        channels = self.mixer['channels']
        any_solo = any(ch.solo for ch in channels)
        any_mic_on = any(ch.mic_on for ch in channels[:2])
        any_pfl = any(ch.pfl for ch in channels)

        for ch in channels:
            in_L = ch.in_port_L.get_array()
            in_R = ch.in_port_R.get_array()

            mute = ch.muted or (any_solo and not ch.solo)
            is_mic = ch.index < 2
            g = ch.gain

            if ch.pfl:
                pfl_buf_L[:] += in_L[:] * g
                pfl_buf_R[:] += in_R[:] * g

            if mute:
                ch.vu_peak = 0.0
                ch._vu_smooth = 0.0
                continue

            factor = g * ch.volume
            goes_to_main = not is_mic or ch.mic_on
            goes_to_headphone = not is_mic

            if goes_to_main:
                master_buf_L[:] += in_L[:] * factor
                master_buf_R[:] += in_R[:] * factor
            if goes_to_headphone:
                headphone_buf_L[:] += in_L[:] * factor
                headphone_buf_R[:] += in_R[:] * factor

            peak = max(np.max(np.abs(in_L)), np.max(np.abs(in_R))) * factor
            ch._vu_smooth = ch._vu_smooth * vu_decay + peak * (1 - vu_decay)
            ch.vu_peak = ch._vu_smooth

        # ── Studio bus: same as main, muted when any mic is on ────────
        if any_mic_on:
            studio_buf_L.fill(0.0)
            studio_buf_R.fill(0.0)
        else:
            studio_buf_L[:] = master_buf_L[:]
            studio_buf_R[:] = master_buf_R[:]

        # ── PFL / Headphone blend ───────────────────────────────────────
        if any_pfl:
            blend = self.mixer['master'].pfl_hp_blend
            hp_gain = blend
            pfl_gain = 1.0 - blend
            if pfl_gain > 0.0:
                headphone_buf_L[:] = headphone_buf_L * hp_gain + pfl_buf_L * pfl_gain
                headphone_buf_R[:] = headphone_buf_R * hp_gain + pfl_buf_R * pfl_gain
            else:
                headphone_buf_L[:] = pfl_buf_L[:]
                headphone_buf_R[:] = pfl_buf_R[:]

        # ── Master VU ──────────────────────────────────────────────────
        master = self.mixer['master']
        peak = max(np.max(np.abs(master_buf_L)), np.max(np.abs(master_buf_R)))
        master._vu_smooth = master._vu_smooth * vu_decay + peak * (1 - vu_decay)
        master.vu_peak = master._vu_smooth

        vol = self.mixer['master'].volume
        if vol < 1.0:
            master_buf_L[:] *= vol
            master_buf_R[:] *= vol

        # ── LUFS metering from external processing input ────────────────
        master._lufs_accum += frames
        if master._lufs_accum >= 512:
            lufs_dt = master._lufs_accum / self.sr
            master._lufs_accum -= 512
            lufs_L = self.lufs_in_L.get_array()
            lufs_R = self.lufs_in_R.get_array()
            ms = (np.mean(lufs_L ** 2) + np.mean(lufs_R ** 2)) / 2.0
            if ms < 1e-12:
                master.lufs_current = -70.0
            else:
                master.lufs_current = -0.691 + 10 * np.log10(ms)
            decay_short = np.exp(-lufs_dt / 4.0)
            decay_long = np.exp(-lufs_dt / 600.0)
            master.lufs_short = master.lufs_short * decay_short + master.lufs_current * (1 - decay_short)
            master.lufs_long = master.lufs_long * decay_long + master.lufs_current * (1 - decay_long)

    def load_track(self, deck_key: str, filepath: str):
        def _load():
            audio = load_audio(filepath, self.sr)
            track = Track(filepath, audio, self.sr)
            self.decks[deck_key].load(track)
        threading.Thread(target=_load, daemon=True).start()

    def load_cart_track(self, cart_key: str, filepath: str):
        def _load():
            audio = load_audio(filepath, self.sr)
            audio = normalize_lufs(audio, self.sr)
            track = Track(filepath, audio, self.sr)
            self.decks[cart_key].load(track)
        threading.Thread(target=_load, daemon=True).start()

    def load_queue(self, deck_key: str, filepath: str):
        def _load():
            audio = load_audio(filepath, self.sr)
            track = Track(filepath, audio, self.sr)
            self.decks[deck_key].load_queue(track)
        threading.Thread(target=_load, daemon=True).start()

    def shutdown(self):
        self.client.deactivate()
        self.client.close()