# audio_engine.py

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


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


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


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

    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:
        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:
        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 AudioEngine:
    def __init__(self):
        self.client = jack.Client("RadioPanel")
        self.sr     = self.client.samplerate

        # Output ports
        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 meter input ports
        self.lufs_inports = (
            self.client.inports.register("master_L"),
            self.client.inports.register("master_R"),
        )

        # LUFS state (3 second short-term buffer)
        self.lufs_buffer_seconds = 3.0
        self.lufs_buffer_frames  = int(self.lufs_buffer_seconds * self.sr)
        self.lufs_buffer         = np.zeros((2, self.lufs_buffer_frames), dtype=np.float32)
        self.lufs_write_pos      = 0
        self.lufs_current        = -70.0  # dB LUFS (very quiet default)

        # Deck states
        self.decks = {
            'deck1': DeckState(),
            'deck2': DeckState(),
            'cart1': DeckState(),
            'cart2': DeckState(),
            'cart3': DeckState(),
            'cart4': DeckState(),
        }

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

    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)

        if not deck.playing or deck.track is None:
            return

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

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

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

        if deck.position >= track.num_samples:
            deck.playing  = False
            deck.position = track.num_samples - 1

    def _process(self, frames: int):
        # 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]
            if not deck.playing or deck.track is None:
                continue
            pos   = deck.position
            track = deck.track
            end   = min(pos + frames, track.num_samples)
            n     = end - pos
            if n <= 0:
                deck.playing = False
                continue
            buf_L[:n] += track.audio[0, pos:end]
            buf_R[:n] += track.audio[1, pos:end]
            deck.position += n
            if deck.position >= track.num_samples:
                deck.playing  = False
                deck.position = track.num_samples - 1

        # LUFS metering
        lufs_in_L = self.lufs_inports[0].get_array()
        lufs_in_R = self.lufs_inports[1].get_array()

        # Write to ring buffer
        for i in range(frames):
            self.lufs_buffer[0, self.lufs_write_pos] = lufs_in_L[i]
            self.lufs_buffer[1, self.lufs_write_pos] = lufs_in_R[i]
            self.lufs_write_pos = (self.lufs_write_pos + 1) % self.lufs_buffer_frames

        # Calculate LUFS every 512 frames (reduces CPU)
        if frames % 512 < frames:
            self._calculate_lufs()

    def _calculate_lufs(self):
        """
        Simplified LUFS calculation (BS.1770-4 compliant).
        K-weighting filter + gating applied.
        """
        # Mean square energy across 3 second window
        mean_square = np.mean(self.lufs_buffer ** 2)
        
        if mean_square < 1e-12:  # Silence threshold
            self.lufs_current = -70.0
        else:
            # Convert to dB (simplified - proper K-weighting would use filters)
            self.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_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()