"""Audio utilities for processing and buffering."""

from __future__ import annotations
from typing import Optional, Tuple, List
import threading, time
from collections import deque
from dataclasses import dataclass
import numpy as np


@dataclass
class AudioChunk:
    samples: np.ndarray
    timestamp: float
    rms_db: float


class EnergyVAD:
    """Energy-based Voice Activity Detection."""
    def __init__(self, threshold_db: float = -35.0):
        self.threshold_db = threshold_db

    def is_speech(self, samples: np.ndarray) -> bool:
        rms = np.sqrt(np.mean(samples ** 2))
        db = 20 * np.log10(rms + 1e-10)
        return db > self.threshold_db

    def rms_db(self, samples: np.ndarray) -> float:
        rms = np.sqrt(np.mean(samples ** 2))
        return 20 * np.log10(rms + 1e-10)


class UtteranceSegmenter:
    """Buffer and segment audio utterances based on energy VAD."""
    def __init__(self, vad: Optional[EnergyVAD] = None, silence_duration_s: float = 0.5,
                 min_duration_s: float = 0.3, sample_rate: int = 16000):
        self.vad = vad or EnergyVAD()
        self.silence_duration_s = silence_duration_s
        self.min_duration_s = min_duration_s
        self.sample_rate = sample_rate
        self._buffer = deque()
        self._last_speech_time = 0.0
        self._speech_started = False
        self._lock = threading.Lock()

    def push(self, samples: np.ndarray) -> List[Tuple[List[float], float]]:
        completed = []
        with self._lock:
            now = time.time()
            is_speech = self.vad.is_speech(samples)
            if is_speech:
                self._last_speech_time = now
                self._speech_started = True
                self._buffer.append(samples)
            else:
                self._buffer.append(samples)
                if self._speech_started and now - self._last_speech_time > self.silence_duration_s:
                    utt_samples = self._extract_buffer()
                    duration = len(utt_samples) / self.sample_rate
                    if duration >= self.min_duration_s:
                        completed.append((utt_samples, duration))
                    self._speech_started = False
                    self._buffer.clear()
        return completed

    def _extract_buffer(self) -> List[float]:
        if not self._buffer:
            return []
        flat = []
        for s in self._buffer:
            flat.extend(s.tolist() if isinstance(s, np.ndarray) else s)
        return flat

    def reset(self) -> None:
        with self._lock:
            self._buffer.clear()
            self._speech_started = False


class AudioBuffer:
    """Thread-safe circular audio buffer."""
    def __init__(self, capacity_s: float = 5.0, sample_rate: int = 16000):
        self.capacity = int(capacity_s * sample_rate)
        self.sample_rate = sample_rate
        self._buffer = deque(maxlen=self.capacity)
        self._lock = threading.Lock()

    def push(self, samples: np.ndarray) -> None:
        with self._lock:
            self._buffer.extend(samples.tolist() if isinstance(samples, np.ndarray) else samples)

    def extract(self, duration_s: Optional[float] = None) -> np.ndarray:
        with self._lock:
            samples = list(self._buffer)
            if duration_s is not None:
                num_samples = int(duration_s * self.sample_rate)
                samples = samples[-num_samples:]
            return np.array(samples, dtype=np.float32)

    def clear(self) -> None:
        with self._lock:
            self._buffer.clear()

    def size(self) -> int:
        with self._lock:
            return len(self._buffer)


def pcm16_to_float32(pcm_bytes: bytes) -> np.ndarray:
    samples = np.frombuffer(pcm_bytes, dtype=np.int16)
    return samples.astype(np.float32) / 32768.0


def float32_to_pcm16(samples: np.ndarray) -> bytes:
    if isinstance(samples, list):
        samples = np.array(samples, dtype=np.float32)
    clipped = np.clip(samples, -1.0, 1.0)
    pcm = (clipped * 32767).astype(np.int16)
    return pcm.tobytes()


def resample_audio(samples: np.ndarray, orig_rate: int, target_rate: int) -> np.ndarray:
    if orig_rate == target_rate:
        return samples
    from scipy import signal
    num_samples = int(len(samples) * target_rate / orig_rate)
    resampled = signal.resample(samples, num_samples)
    return resampled.astype(np.float32)


def calculate_rms_db(samples: np.ndarray) -> float:
    rms = np.sqrt(np.mean(samples ** 2))
    return 20 * np.log10(rms + 1e-10)


class VoskSTT:
    """Offline STT backend using Vosk (low-latency alternative to Whisper; Issue #503)."""

    def __init__(self, model_path: str = "/models/vosk/vosk-model-en-us-0.22",
                 sample_rate: int = 16000):
        self.model_path = model_path
        self.sample_rate = sample_rate
        self._recognizer = None
        self._model = None
        self._json = None

    def load(self) -> bool:
        try:
            from vosk import Model, KaldiRecognizer
            import json as _json
            self._json = _json
            self._model = Model(self.model_path)
            self._recognizer = KaldiRecognizer(self._model, self.sample_rate)
            self._recognizer.SetWords(True)
            return True
        except Exception as e:
            print(f"VoskSTT load failed: {e}")
            return False

    def transcribe(self, audio_samples) -> str:
        if self._recognizer is None:
            return ""
        try:
            import numpy as np
            if not isinstance(audio_samples, np.ndarray):
                audio_samples = np.array(audio_samples, dtype=np.float32)
            pcm = float32_to_pcm16(audio_samples)
            self._recognizer.AcceptWaveform(pcm)
            result = self._json.loads(self._recognizer.FinalResult())
            return result.get("text", "").strip()
        except Exception as e:
            print(f"VoskSTT transcribe error: {e}")
            return ""
