"""STT-Kern: Whisper-Engine, Audio-Hilfen und Äußerungs-Segmentierung (VAD).

Wird von stt_service.py (HTTP/WebSocket/Wyoming) und stt_cli.py (lokaler Modus)
gemeinsam genutzt. Keine Ausgabe per print, keine sys.exit() — Fehler werden als
Exceptions gemeldet, Ergebnisse als Datenobjekte zurückgegeben.

Grundsätze (Lehren aus legacy/stt.py und my_voice_assistant_v3):
- Modell EINMAL laden und halten; bei GPU-Problemen robust auf CPU ausweichen.
- Äußerungsende per Silero-VAD (Stille nach Sprache) statt Stoppwort.
- Genau eine Transkription pro Äußerung, ohne Kontext-Verkettung
  (condition_on_previous_text=False) und mit Halluzinationsfilter.
"""

from __future__ import annotations

import io
import logging
import os
import re
import threading
import time
from collections import deque
from dataclasses import asdict, dataclass, field
from math import gcd
from typing import Callable

import numpy as np

log = logging.getLogger("stt")

__version__ = "0.1.0"

SAMPLE_RATE = 16000  # Whisper arbeitet intern mit 16 kHz mono float32


class AudioDecodeError(ValueError):
    """Audio konnte nicht dekodiert werden (defekte Datei, unbekanntes Format)."""


# === Audio-Hilfen =============================================================

def pcm16_to_float(data: bytes, channels: int = 1, channel: int | None = None) -> np.ndarray:
    """Rohes s16le-PCM → float32 mono.

    channels > 1: `channel` wählt einen Kanal (z. B. 1 = rechter ASR-Kanal des
    ReSpeaker XVF3800), sonst werden alle Kanäle gemittelt. Unvollständige
    Frames am Ende werden verworfen.
    """
    frame = 2 * channels
    usable = len(data) - (len(data) % frame)
    audio = np.frombuffer(data[:usable], dtype="<i2").astype(np.float32) / 32768.0
    if channels > 1:
        audio = audio.reshape(-1, channels)
        if channel is not None:
            if not 0 <= channel < channels:
                raise ValueError(f"channel={channel} ungültig bei {channels} Kanälen")
            audio = audio[:, channel]
        else:
            audio = audio.mean(axis=1)
    return np.ascontiguousarray(audio, dtype=np.float32)


def float_to_pcm16(audio: np.ndarray) -> bytes:
    return (np.clip(audio, -1.0, 1.0) * 32767.0).astype("<i2").tobytes()


def resample(audio: np.ndarray, src_rate: int, dst_rate: int = SAMPLE_RATE) -> np.ndarray:
    """Resampelt ein komplettes Signal (nicht für Streaming-Stückchen gedacht)."""
    if src_rate == dst_rate or audio.size == 0:
        return audio.astype(np.float32, copy=False)
    from scipy.signal import resample_poly

    g = gcd(src_rate, dst_rate)
    return resample_poly(audio, dst_rate // g, src_rate // g).astype(np.float32)


class StreamResampler:
    """Resampelt fortlaufende Audio-Stückchen ohne Artefakte an den Stückgrenzen.

    Nutzt soxr.ResampleStream (zustandsbehaftet); ohne soxr Fallback auf
    resample_poly pro Stück (kleine Randartefakte, für ASR unkritisch).
    """

    def __init__(self, src_rate: int, dst_rate: int = SAMPLE_RATE):
        self.src_rate = src_rate
        self.dst_rate = dst_rate
        self._stream = None
        if src_rate != dst_rate:
            try:
                import soxr

                self._stream = soxr.ResampleStream(src_rate, dst_rate, 1, dtype="float32")
            except ImportError:
                self._stream = None

    def process(self, audio: np.ndarray, last: bool = False) -> np.ndarray:
        if self.src_rate == self.dst_rate:
            return audio
        if self._stream is not None:
            return self._stream.resample_chunk(audio.astype(np.float32, copy=False), last=last)
        return resample(audio, self.src_rate, self.dst_rate)


def decode_audio_bytes(data: bytes) -> np.ndarray:
    """Dekodiert beliebige Audio-Container (WAV jeder Bittiefe/Kanalzahl, MP3, OGG,
    WebM, M4A …) via PyAV nach 16 kHz mono float32."""
    if not data:
        raise AudioDecodeError("Leere Audiodaten")
    from faster_whisper import decode_audio

    try:
        return decode_audio(io.BytesIO(data), sampling_rate=SAMPLE_RATE)
    except Exception as exc:  # PyAV wirft diverse Typen
        raise AudioDecodeError(f"Audio nicht dekodierbar: {exc}") from exc


# === Ergebnis-Datenobjekte ====================================================

@dataclass
class Segment:
    start: float
    end: float
    text: str
    avg_logprob: float = 0.0
    no_speech_prob: float = 0.0


@dataclass
class Transcript:
    text: str
    language: str | None
    language_probability: float
    duration: float           # Audiodauer in Sekunden
    processing_s: float       # Rechenzeit in Sekunden
    segments: list[Segment] = field(default_factory=list)
    dropped_segments: int = 0  # vom Halluzinationsfilter verworfen

    def to_dict(self, with_segments: bool = True) -> dict:
        data = asdict(self)
        if not with_segments:
            data.pop("segments")
        return data


# === Halluzinationsfilter =====================================================

# Whisper erzeugt in Stille/Rauschen typische Floskeln aus seinen Trainingsdaten
# (Untertitel). Gefiltert wird nur, wenn ein Segment KOMPLETT aus so einer Floskel
# besteht — echte Sätze, die solche Wörter enthalten, bleiben erhalten.
_HALLUCINATION_PATTERNS = [
    r"untertitel(ung)? (im auftrag )?(des|der|von|by) .*",
    r"untertitel (der|von) amara\.org( community)?",
    r"(copyright|©) ?(wdr|swr|zdf|ard|ndr|br|mdr)( \d{4})?",
    r"(swr|wdr|zdf|ndr|mdr) \d{4}",
    r"vielen dank (fürs|für's|für das) (zuschauen|zusehen|zuhören)",
    r"(thanks|thank you) for watching",
    r"subtitles by .*",
    r"sous-titr(es|age) (réalisés? )?par .*",
]
_HALLUCINATION_RE = re.compile(
    "^(?:" + "|".join(_HALLUCINATION_PATTERNS) + ")$", re.IGNORECASE
)
_PUNCT_RE = re.compile(r"[^\w\s©'.-]+|[.!?…]+$")


def is_hallucination(seg: Segment) -> bool:
    norm = _PUNCT_RE.sub("", seg.text.strip().lower()).strip(" .")
    if not norm:
        return True
    if _HALLUCINATION_RE.match(norm):
        return True
    # Whisper selbst unsicher UND Stille wahrscheinlich → verwerfen
    return seg.no_speech_prob > 0.6 and seg.avg_logprob < -0.8


# === Engine ===================================================================

def _env(name: str, default, cast=str):
    raw = os.environ.get(name)
    if raw is None or raw == "":
        return default
    if cast is bool:
        return raw.strip().lower() in ("1", "true", "yes", "on")
    return cast(raw)


@dataclass
class EngineConfig:
    model: str = "large-v3-turbo"  # tiny|base|small|medium|large-v3|large-v3-turbo|Pfad
    device: str = "auto"           # auto|cuda|cpu
    compute_type: str = "auto"     # auto → float16 (cuda) / int8 (cpu)
    device_index: int = 0          # CUDA-Index; besser GPU per CUDA_VISIBLE_DEVICES=UUID pinnen
    cpu_model: str = "small"       # Modell, wenn (automatisch) auf CPU ausgewichen wird
    cpu_threads: int = 0           # 0 = ctranslate2-Default
    language: str | None = "de"    # Default-Sprache; None/"auto" = automatische Erkennung
    beam_size: int = 5
    initial_prompt: str | None = None
    download_root: str | None = None

    @classmethod
    def from_env(cls) -> "EngineConfig":
        d = cls()
        language = _env("STT_LANGUAGE", d.language)
        return cls(
            model=_env("STT_MODEL", d.model),
            device=_env("STT_DEVICE", d.device),
            compute_type=_env("STT_COMPUTE_TYPE", d.compute_type),
            device_index=_env("STT_DEVICE_INDEX", d.device_index, int),
            cpu_model=_env("STT_CPU_MODEL", d.cpu_model),
            cpu_threads=_env("STT_CPU_THREADS", d.cpu_threads, int),
            language=None if language in (None, "auto") else language,
            beam_size=_env("STT_BEAM_SIZE", d.beam_size, int),
            initial_prompt=_env("STT_INITIAL_PROMPT", d.initial_prompt),
            download_root=_env("STT_DOWNLOAD_ROOT", d.download_root),
        )


def cuda_device_count() -> int:
    try:
        import ctranslate2

        return ctranslate2.get_cuda_device_count()
    except Exception:
        return 0


class SttEngine:
    """Thread-sichere Whisper-Engine. Modell wird beim ersten Bedarf (oder per
    load()) genau einmal geladen; parallele Aufrufe werden serialisiert."""

    def __init__(self, config: EngineConfig | None = None):
        self.config = config or EngineConfig()
        self._model = None
        self._load_lock = threading.Lock()
        self._run_lock = threading.Lock()
        self.status = "idle"  # idle|loading|ready|error
        self.error: str | None = None
        self.active_model: str | None = None
        self.active_device: str | None = None
        self.active_compute_type: str | None = None
        self.load_seconds: float | None = None

    # --- Laden ----------------------------------------------------------------

    def _load_attempts(self) -> list[tuple[str, str, str]]:
        cfg = self.config
        attempts: list[tuple[str, str, str]] = []
        if cfg.device in ("auto", "cuda"):
            if cuda_device_count() > 0:
                ct = "float16" if cfg.compute_type == "auto" else cfg.compute_type
                attempts.append(("cuda", cfg.model, ct))
            else:
                log.warning("Keine CUDA-GPU gefunden — nutze CPU")
        cpu_model = cfg.model if cfg.device == "cpu" else cfg.cpu_model
        cpu_ct = "int8" if cfg.compute_type == "auto" or cfg.device != "cpu" else cfg.compute_type
        attempts.append(("cpu", cpu_model, cpu_ct))
        return attempts

    def load(self):
        with self._load_lock:
            if self._model is not None:
                return self._model
            from faster_whisper import WhisperModel

            self.status = "loading"
            t0 = time.perf_counter()
            last_exc: Exception | None = None
            for device, model_name, compute_type in self._load_attempts():
                try:
                    log.info("Lade Whisper-Modell %s (%s, %s)", model_name, device, compute_type)
                    kwargs = dict(
                        device=device,
                        compute_type=compute_type,
                        cpu_threads=self.config.cpu_threads,
                        download_root=self.config.download_root,
                    )
                    if device == "cuda":
                        kwargs["device_index"] = self.config.device_index
                    self._model = WhisperModel(model_name, **kwargs)
                    self.active_model, self.active_device = model_name, device
                    self.active_compute_type = compute_type
                    break
                except Exception as exc:
                    last_exc = exc
                    log.warning("Laden auf %s fehlgeschlagen: %s", device, exc)
            if self._model is None:
                self.status, self.error = "error", str(last_exc)
                raise RuntimeError(f"Whisper-Modell nicht ladbar: {last_exc}") from last_exc
            self.load_seconds = time.perf_counter() - t0
            self.status, self.error = "ready", None
            log.info("Modell bereit nach %.1fs", self.load_seconds)
            return self._model

    def warmup(self) -> None:
        """Lädt das Modell und rechnet einmal durch (CUDA-Kernel, Decoder)."""
        self.load()
        noise = (np.random.default_rng(0).standard_normal(SAMPLE_RATE) * 0.01).astype(np.float32)
        self.transcribe(noise, vad_filter=False, beam_size=1)

    # --- Transkription ---------------------------------------------------------

    def resolve_language(self, language: str | None) -> str | None:
        """None → Default aus Config; 'auto' → Erkennung; 'de-DE' → 'de'."""
        if language is None:
            return self.config.language
        language = language.strip().lower()
        if language in ("", "auto"):
            return None
        return language.split("-")[0].split("_")[0]

    def transcribe(
        self,
        audio: np.ndarray,
        *,
        language: str | None = None,
        initial_prompt: str | None = None,
        hotwords: str | None = None,
        beam_size: int | None = None,
        vad_filter: bool = True,
    ) -> Transcript:
        """Transkribiert 16-kHz-mono-float32-Audio. Blockiert — in Diensten in einem
        Worker-Thread aufrufen."""
        duration = float(audio.size) / SAMPLE_RATE
        lang = self.resolve_language(language)
        if audio.size < SAMPLE_RATE // 10:  # < 100 ms: nichts Sinnvolles
            return Transcript("", lang, 0.0, duration, 0.0)

        model = self.load()
        t0 = time.perf_counter()
        with self._run_lock:
            segments, info = model.transcribe(
                audio.astype(np.float32, copy=False),
                language=lang,
                beam_size=beam_size or self.config.beam_size,
                vad_filter=vad_filter,
                vad_parameters={"min_silence_duration_ms": 500, "speech_pad_ms": 200},
                condition_on_previous_text=False,
                initial_prompt=initial_prompt or self.config.initial_prompt,
                hotwords=hotwords,
            )
            raw = list(segments)  # Generator: die eigentliche Arbeit passiert hier

        kept: list[Segment] = []
        dropped = 0
        for s in raw:
            seg = Segment(round(s.start, 2), round(s.end, 2), s.text.strip(),
                          round(s.avg_logprob, 3), round(s.no_speech_prob, 3))
            if is_hallucination(seg):
                dropped += 1
                log.debug("Halluzination verworfen: %r", seg.text)
            else:
                kept.append(seg)
        text = " ".join(s.text for s in kept).strip()
        return Transcript(
            text=text,
            language=info.language,
            language_probability=round(float(info.language_probability), 3),
            duration=round(duration, 2),
            processing_s=round(time.perf_counter() - t0, 3),
            segments=kept,
            dropped_segments=dropped,
        )


# === Äußerungs-Segmentierung (Streaming-VAD) ==================================

VAD_WINDOW = 512  # Samples pro Silero-Fenster bei 16 kHz (= 32 ms)
_WINDOW_MS = VAD_WINDOW * 1000 / SAMPLE_RATE


class SileroProb:
    """Sprachwahrscheinlichkeit je 512-Sample-Fenster mit dem in faster-whisper
    mitgelieferten Silero-VAD (ONNX, CPU, kein Download).

    Das Modell leitet den Kontext eines Fensters aus dem vorherigen Fenster im
    selben Aufruf ab. Damit das über Aufrufgrenzen stimmt, wird das letzte Fenster
    des Vorgängeraufrufs vorangestellt und dessen Ergebnis verworfen.
    """

    def __init__(self):
        from faster_whisper.vad import get_vad_model

        self._model = get_vad_model()
        self._prev = np.zeros(VAD_WINDOW, dtype=np.float32)

    def __call__(self, windows: np.ndarray) -> np.ndarray:
        audio = np.concatenate([self._prev, windows])
        self._prev = windows[-VAD_WINDOW:].copy()
        return np.asarray(self._model(audio)).reshape(-1)[1:]

    def reset(self) -> None:
        self._prev[:] = 0


@dataclass
class VadConfig:
    threshold: float = 0.5        # Sprachwahrscheinlichkeit für "Sprache"
    min_speech_ms: int = 200      # so lange Sprache nötig, bis eine Äußerung beginnt
    silence_ms: int = 700         # so lange Stille beendet eine Äußerung
    pre_roll_ms: int = 300        # Audio vor Sprachbeginn mitnehmen (Anlaute)
    max_utterance_s: float = 30.0  # harte Obergrenze (Whisper-Fenster)

    @classmethod
    def from_env(cls) -> "VadConfig":
        d = cls()
        return cls(
            threshold=_env("STT_VAD_THRESHOLD", d.threshold, float),
            min_speech_ms=_env("STT_VAD_MIN_SPEECH_MS", d.min_speech_ms, int),
            silence_ms=_env("STT_VAD_SILENCE_MS", d.silence_ms, int),
            pre_roll_ms=_env("STT_VAD_PRE_ROLL_MS", d.pre_roll_ms, int),
            max_utterance_s=_env("STT_VAD_MAX_UTTERANCE_S", d.max_utterance_s, float),
        )


@dataclass
class VadEvent:
    kind: str                      # "speech_start" | "utterance"
    time_s: float                  # Stream-Zeit des Ereignisses
    audio: np.ndarray | None = None
    reason: str | None = None      # utterance: "silence" | "max_length" | "flush"


class UtteranceSegmenter:
    """Zerlegt einen fortlaufenden 16-kHz-Stream in Äußerungen.

    feed() liefert Ereignisse, sobald Sprache beginnt bzw. eine Äußerung durch
    Stille, Maximallänge oder flush() endet. Rein rechnerisch, kein I/O.
    """

    def __init__(self, config: VadConfig | None = None,
                 prob_fn: Callable[[np.ndarray], np.ndarray] | None = None):
        self.config = config or VadConfig()
        self._prob_fn = prob_fn if prob_fn is not None else SileroProb()
        pre_windows = int(np.ceil((self.config.pre_roll_ms + self.config.min_speech_ms) / _WINDOW_MS))
        self._pre: deque[np.ndarray] = deque(maxlen=max(pre_windows, 1))
        self._pending = np.zeros(0, dtype=np.float32)
        self._speech: list[np.ndarray] | None = None
        self._voiced = 0
        self._silent = 0
        self._samples = 0

    @property
    def in_speech(self) -> bool:
        return self._speech is not None

    def reset(self) -> None:
        self._pre.clear()
        self._pending = np.zeros(0, dtype=np.float32)
        self._speech = None
        self._voiced = self._silent = 0
        if hasattr(self._prob_fn, "reset"):
            self._prob_fn.reset()

    def _now(self) -> float:
        return self._samples / SAMPLE_RATE

    def _finish(self, reason: str) -> VadEvent:
        audio = np.concatenate(self._speech) if self._speech else np.zeros(0, np.float32)
        self._speech = None
        self._voiced = self._silent = 0
        self._pre.clear()
        return VadEvent("utterance", self._now(), audio=audio, reason=reason)

    def feed(self, audio: np.ndarray) -> list[VadEvent]:
        cfg = self.config
        buf = np.concatenate([self._pending, audio.astype(np.float32, copy=False)])
        n = buf.size // VAD_WINDOW
        self._pending = buf[n * VAD_WINDOW:]
        if n == 0:
            return []
        windows = buf[: n * VAD_WINDOW]
        probs = self._prob_fn(windows)

        min_speech = max(1, int(cfg.min_speech_ms / _WINDOW_MS))
        end_silence = max(1, int(cfg.silence_ms / _WINDOW_MS))
        max_windows = int(cfg.max_utterance_s * 1000 / _WINDOW_MS)
        neg_threshold = max(cfg.threshold - 0.15, 0.01)

        events: list[VadEvent] = []
        for i in range(n):
            win = windows[i * VAD_WINDOW:(i + 1) * VAD_WINDOW]
            p = float(probs[i])
            self._samples += VAD_WINDOW
            if self._speech is None:
                self._pre.append(win)
                self._voiced = self._voiced + 1 if p >= cfg.threshold else 0
                if self._voiced >= min_speech:
                    self._speech = list(self._pre)
                    self._pre.clear()
                    self._silent = 0
                    events.append(VadEvent("speech_start", self._now()))
            else:
                self._speech.append(win)
                self._silent = self._silent + 1 if p < neg_threshold else 0
                if self._silent >= end_silence:
                    events.append(self._finish("silence"))
                elif len(self._speech) >= max_windows:
                    events.append(self._finish("max_length"))
        return events

    def flush(self) -> list[VadEvent]:
        """Beendet eine laufende Äußerung (z. B. Stream-Ende)."""
        if self._speech is None:
            self._pending = np.zeros(0, dtype=np.float32)
            return []
        if self._pending.size:
            self._speech.append(self._pending)
            self._pending = np.zeros(0, dtype=np.float32)
        return [self._finish("flush")]