# 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 = [
    "Turntable 1", "Turntable 2", "Mic", "Aux 1",
    "Aux 2", "Aux 3", "Aux 4", "Aux 5",
]


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 ExternalInputMonitor:
    def __init__(self, buffer_size=600):
        self.buffer = np.zeros(buffer_size, dtype=np.float32)
        self.write_pos = 0
        self.buffer_size = buffer_size
        self.monitoring = False


class MixerChannel:
    def __init__(self, name: str, index: int):
        self.name     = name
        self.index    = index
        self.volume   = 1.0
        self.muted    = False
        self.solo     = False
        self.pfl     = False
        self.vu_peak  = 0.0
        self.in_port_L = None
        self.in_port_R = None


class Compressor:
    def __init__(self, sr: int):
        self.sr = sr
        self.threshold_db = -20.0
        self.ratio = 4.0
        self.attack_ms = 5.0
        self.release_ms = 100.0
        self.makeup_gain_db = 0.0
        self.bypassed = False
        self._envelope = 0.0
        self._attack_coeff = self._time_constant(self.attack_ms)
        self._release_coeff = self._time_constant(self.release_ms)

    def _time_constant(self, ms: float) -> float:
        return np.exp(-1.0 / (ms * self.sr / 1000.0))

    def set_attack(self, ms: float):
        self.attack_ms = ms
        self._attack_coeff = self._time_constant(ms)

    def set_release(self, ms: float):
        self.release_ms = ms
        self._release_coeff = self._time_constant(ms)

    def process(self, audio: np.ndarray):
        if self.bypassed:
            return
        threshold_linear = 10 ** (self.threshold_db / 20.0)
        makeup_linear = 10 ** (self.makeup_gain_db / 20.0)
        for ch in range(audio.shape[0]):
            for i in range(audio.shape[1]):
                sample = audio[ch, i]
                level = abs(sample)
                if level > self._envelope:
                    self._envelope += (1 - self._attack_coeff) * (level - self._envelope)
                else:
                    self._envelope += (1 - self._release_coeff) * (level - self._envelope)
                if self._envelope > threshold_linear:
                    env_db = 20 * np.log10(max(self._envelope, 1e-10))
                    gain_db = (self.threshold_db - env_db) * (1 - 1.0 / self.ratio)
                    gain = 10 ** (gain_db / 20.0)
                else:
                    gain = 1.0
                audio[ch, i] = sample * gain * makeup_linear


class Limiter:
    def __init__(self):
        self.ceiling_db = -0.1
        self.bypassed = False

    def process(self, audio: np.ndarray):
        if self.bypassed:
            return
        ceiling = 10 ** (self.ceiling_db / 20.0)
        np.clip(audio, -ceiling, ceiling, out=audio)


class MasterBus:
    def __init__(self, sr: int):
        self.compressor = Compressor(sr)
        self.limiter = Limiter()
        self.vu_peak = 0.0
        self.lufs_current = -70.0
        self.volume = 1.0
        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"),
            ),
        }

        # External input monitors (deck3, deck4 — pass-through with waveform)
        self.ext_monitors = {}
        self.ext_ports = {}
        for ext_key, label in [('deck3', 'deck3'), ('deck4', 'deck4')]:
            in_L = self.client.inports.register(f"{label}_in_L")
            in_R = self.client.inports.register(f"{label}_in_R")
            out_L = self.client.outports.register(f"{label}_out_L")
            out_R = self.client.outports.register(f"{label}_out_R")
            self.ext_ports[ext_key] = (in_L, in_R, out_L, out_R)
            self.ext_monitors[ext_key] = ExternalInputMonitor()

        # 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"),
        )

        # 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_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 _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

        # ── External input monitors (deck3, deck4) ─────────────────────
        gate_linear = 10 ** (-46.0 / 20.0)
        for ext_key in ('deck3', 'deck4'):
            in_L, in_R, out_L, out_R = self.ext_ports[ext_key]
            in_buf_L = in_L.get_array()
            in_buf_R = in_R.get_array()
            out_buf_L = out_L.get_array()
            out_buf_R = out_R.get_array()

            out_buf_L[:] = in_buf_L[:]
            out_buf_R[:] = in_buf_R[:]

            monitor = self.ext_monitors[ext_key]
            if not monitor.monitoring:
                monitor.buffer[:] = 0.0
                continue

            rms = np.sqrt(np.mean((in_buf_L ** 2 + in_buf_R ** 2) / 2.0))
            if rms > gate_linear:
                db = 20 * np.log10(max(rms, 1e-10))
                display = max(0.0, (db + 46.0) / 46.0)
                display = min(1.0, display)
            else:
                display = 0.0

            monitor.buffer[monitor.write_pos] = display
            monitor.write_pos = (monitor.write_pos + 1) % monitor.buffer_size

        # ── Mixer ───────────────────────────────────────────────────────
        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)

        any_solo = any(ch.solo for ch in self.mixer['channels'])

        # PFL bus
        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)

        for ch in self.mixer['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)

            # PFL: pre-fader, pre-mute signal to PFL bus
            if ch.pfl:
                pfl_buf_L[:] += in_L[:]
                pfl_buf_R[:] += in_R[:]

            if mute:
                ch.vu_peak = 0.0
                continue

            vol = ch.volume
            peak = 0.0
            for i in range(frames):
                s_L = in_L[i] * vol
                s_R = in_R[i] * vol
                master_buf_L[i] += s_L
                master_buf_R[i] += s_R
                peak = max(peak, abs(s_L), abs(s_R))

            ch.vu_peak = peak

        # ── Master DSP ──────────────────────────────────────────────────
        master_audio = np.array([master_buf_L, master_buf_R])

        self.mixer['master'].compressor.process(master_audio)
        self.mixer['master'].limiter.process(master_audio)

        master_buf_L[:] = master_audio[0]
        master_buf_R[:] = master_audio[1]

        master_peak = max(np.max(np.abs(master_audio[0])), np.max(np.abs(master_audio[1])))
        self.mixer['master'].vu_peak = master_peak

        master = self.mixer['master']
        master._lufs_accum += frames
        if master._lufs_accum >= 512:
            master._lufs_accum -= 512
            self._calculate_lufs(master_audio)

    def _calculate_lufs(self, audio: np.ndarray):
        mean_square = np.mean(audio ** 2)
        if mean_square < 1e-12:
            self.mixer['master'].lufs_current = -70.0
        else:
            self.mixer['master'].lufs_current = -0.691 + 10 * np.log10(mean_square)

    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()