#!/usr/bin/env python3
"""
ARIA Voxtral Bridge — Streaming-STT via Voxtral-Mini-4B-Realtime-2602 (vLLM).

Zwilling der whisper-Bridge, aber das Transkribieren macht NICHT faster-whisper
im selben Prozess, sondern der separate vLLM-Realtime-Server (Container
`voxtral-vllm`) ueber dessen WebSocket-API `/v1/realtime`. Diese Bridge ist
reine Glue:

  App ──(RVS: stt_stream_start / stt_audio_chunk / stt_stream_end)──▶ diese Bridge
  diese Bridge ──(WS /v1/realtime: PCM16-b64 append)──▶ voxtral-vllm
  voxtral-vllm ──(transcription.delta / transcription.done)──▶ diese Bridge
  diese Bridge ──(RVS: stt_partial / stt_endpoint / stt_stream_done)──▶ App/aria-bridge

Das RVS-Wire-Protokoll ist IDENTISCH zur whisper-Bridge (drop-in). Das
Endpointing (wann hat der User aufgehoert zu sprechen) macht diese Bridge
selbst — mit demselben ADAPTIVEN Rausch-Boden-Endpointer wie whisper (Voxtral
Realtime liefert laut vLLM-Doku keine eigene VAD/„speaker done"-Semantik, nur
transcription.delta/.done). Die akustische Energie messen wir auf unserer
eigenen PCM-Kopie, die semantische Stagnation am Delta-Textwachstum.

⚠️  HARDWARE: Voxtral-Mini-4B-Realtime-2602 braucht >=16 GB VRAM (BF16). Auf der
    RTX 3060 (12 GB) laeuft es NICHT — erst auf der 24-GB-Karte. Bis dahin
    bleibt die whisper-Bridge aktiv (Profil-gesteuert im docker-compose).

⚠️  VERIFY-ON-FIRST-RUN: Die exakten vLLM-Realtime-FRAME-Namen (Audio-Append,
    Delta/Done-Event-Typen) sind unten als Konstanten gebuendelt und nach dem
    OpenAI-Realtime-Schema modelliert. Gegen das offizielle vLLM-Realtime-
    Client-Beispiel pruefen und ggf. anpassen — sie stehen bewusst an EINER
    Stelle. Response-Handling ist defensiv (mehrere moegliche Feldnamen).

Env:
  RVS_HOST, RVS_PORT, RVS_TLS, RVS_TLS_FALLBACK, RVS_TOKEN
  VOXTRAL_VLLM_URL     Default: ws://voxtral-vllm:8000/v1/realtime
  VOXTRAL_MODEL        Default: mistralai/Voxtral-Mini-4B-Realtime-2602
  VOXTRAL_LANGUAGE     Default: de
"""
import asyncio
import base64
import json
import logging
import os
import time
from dataclasses import dataclass, field
from typing import Optional

import numpy as np
import websockets

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger("voxtral-bridge")

RVS_HOST = os.getenv("RVS_HOST", "").strip()
RVS_PORT = int(os.getenv("RVS_PORT", "443"))
RVS_TLS = os.getenv("RVS_TLS", "true").lower() == "true"
RVS_TLS_FALLBACK = os.getenv("RVS_TLS_FALLBACK", "true").lower() == "true"
RVS_TOKEN = os.getenv("RVS_TOKEN", "").strip()

VOXTRAL_VLLM_URL = os.getenv("VOXTRAL_VLLM_URL", "ws://voxtral-vllm:8000/v1/realtime")
VOXTRAL_MODEL = os.getenv("VOXTRAL_MODEL", "mistralai/Voxtral-Mini-4B-Realtime-2602")
VOXTRAL_LANGUAGE = os.getenv("VOXTRAL_LANGUAGE", "de")

# ── vLLM-Realtime-Frames — HIER anpassen falls das Client-Beispiel abweicht ──
# Senderichtung (wir → vLLM): PCM16-16kHz-mono base64 anhaengen + committen.
VLLM_SEND_APPEND = "input_audio_buffer.append"   # {"type":..., "audio": "<b64>"}
VLLM_SEND_COMMIT = "input_audio_buffer.commit"   # Buffer abschliessen
VLLM_AUDIO_FIELD = "audio"
# Empfangsrichtung (vLLM → wir): inkrementeller Text + final. Defensiv geprueft.
VLLM_DELTA_SUFFIXES = ("transcription.delta",)   # msg["type"] endet hierauf
VLLM_DONE_SUFFIXES = ("transcription.done", "transcription.completed")
VLLM_DELTA_FIELDS = ("delta", "text", "transcription")   # eins davon traegt den Text

# ── Streaming-/Endpointing-Parameter (analog whisper-Bridge) ──
STREAM_DEFAULT_ENDPOINT_MS = 2400
STREAM_DEFAULT_HARD_CAP_MS = 60000
STREAM_MIN_AUDIO_MS = 600
STREAM_SESSION_TTL_S = 120
STREAM_ENERGY_WINDOW_MS = 300
STREAM_SEMANTIC_BACKUP_FACTOR = 2.0
# Adaptiver Voice-Schwellwert (siehe whisper-Bridge M0.1): Grenze relativ zum
# gemessenen Rausch-Boden statt fix — schneidet leises Sprechen nicht ab.
STREAM_VOICE_FACTOR = 2.5
STREAM_VOICE_RMS_MIN = 0.005
STREAM_VOICE_RMS_MAX = 0.020


def pcm_s16le_to_float32(data: bytes) -> np.ndarray:
    if not data:
        return np.zeros(0, dtype=np.float32)
    arr = np.frombuffer(data, dtype=np.int16).astype(np.float32) / 32768.0
    return arr


async def _send(ws, mtype: str, payload: dict) -> None:
    try:
        await ws.send(json.dumps({
            "type": mtype,
            "payload": payload,
            "timestamp": int(time.time() * 1000),
        }))
    except Exception as e:
        logger.warning("RVS-Send fehlgeschlagen (%s): %s", mtype, e)


@dataclass
class StreamSession:
    request_id: str
    audio_request_id: str
    language: str
    endpoint_ms: int
    hard_cap_ms: int
    voice: str = ""
    speed: float = 1.0
    interrupted: bool = False
    location: Optional[dict] = None
    sample_rate: int = 16000
    voice_factor: float = STREAM_VOICE_FACTOR
    voice_rms_min: float = STREAM_VOICE_RMS_MIN
    voice_rms_max: float = STREAM_VOICE_RMS_MAX
    pcm_buffer: bytearray = field(default_factory=bytearray)
    started_at: float = field(default_factory=time.time)
    last_chunk_at: float = field(default_factory=time.time)
    last_partial: str = ""
    last_growth_at: float = 0.0
    last_voice_at: float = 0.0
    noise_floor: float = 0.0
    closed: bool = False
    endpoint_sent: bool = False
    # vLLM-Realtime-Session
    vllm_ws: object = None
    vllm_reader: object = None


class SessionManager:
    def __init__(self) -> None:
        self._sessions: dict[str, StreamSession] = {}
        self._ws = None  # RVS

    def attach_ws(self, ws) -> None:
        self._ws = ws

    async def start_session(self, payload: dict) -> Optional[StreamSession]:
        request_id = (payload.get("requestId") or "").strip()
        if not request_id:
            logger.warning("stt_stream_start ohne requestId — ignoriert")
            return None
        try:
            endpoint_ms = int(payload.get("endpointMs") or STREAM_DEFAULT_ENDPOINT_MS)
        except (TypeError, ValueError):
            endpoint_ms = STREAM_DEFAULT_ENDPOINT_MS
        try:
            hard_cap_ms = int(payload.get("hardCapMs") or STREAM_DEFAULT_HARD_CAP_MS)
        except (TypeError, ValueError):
            hard_cap_ms = STREAM_DEFAULT_HARD_CAP_MS
        try:
            voice_factor = float(payload.get("voiceFactor") or STREAM_VOICE_FACTOR)
        except (TypeError, ValueError):
            voice_factor = STREAM_VOICE_FACTOR
        sess = StreamSession(
            request_id=request_id,
            audio_request_id=payload.get("audioRequestId", "") or "",
            language=payload.get("language") or VOXTRAL_LANGUAGE,
            endpoint_ms=endpoint_ms,
            hard_cap_ms=hard_cap_ms,
            voice=payload.get("voice", "") or "",
            speed=float(payload.get("speed") or 1.0),
            voice_factor=voice_factor,
            interrupted=bool(payload.get("interrupted", False)),
            location=payload.get("location") or None,
            sample_rate=int(payload.get("sampleRate") or 16000),
        )
        # vLLM-Realtime-Session oeffnen + Reader starten.
        try:
            sess.vllm_ws = await websockets.connect(VOXTRAL_VLLM_URL, max_size=8 * 1024 * 1024)
            await self._vllm_configure(sess)
            sess.vllm_reader = asyncio.create_task(self._vllm_read_loop(sess))
        except Exception as e:
            logger.exception("Stream %s: vLLM-Realtime-Connect fehlgeschlagen: %s",
                             request_id[:8], e)
            # ohne Backend keine Transkription → sofort leeres Endpoint melden
            self._sessions[request_id] = sess
            await self._finalize(sess, reason="vllm_unavailable")
            return None
        self._sessions[request_id] = sess
        logger.info("Voxtral-Session offen: id=%s lang=%s endpointMs=%d",
                    request_id[:8], sess.language, sess.endpoint_ms)
        return sess

    async def _vllm_configure(self, sess: StreamSession) -> None:
        """Optionale Session-Konfig an vLLM (Modell/Sprache/temperature=0).
        VERIFY: exaktes session.update-Schema gegen vLLM-Realtime-Beispiel.
        Best-effort — Fehler hier sind nicht fatal."""
        try:
            await sess.vllm_ws.send(json.dumps({
                "type": "session.update",
                "session": {
                    "model": VOXTRAL_MODEL,
                    "language": sess.language,
                    "temperature": 0.0,
                    "input_audio_format": "pcm16",
                },
            }))
        except Exception:
            pass

    async def _vllm_read_loop(self, sess: StreamSession) -> None:
        """Liest transcription.delta/.done vom vLLM-Realtime-Server."""
        ws = sess.vllm_ws
        try:
            async for raw in ws:
                try:
                    msg = json.loads(raw)
                except Exception:
                    continue
                mtype = str(msg.get("type", ""))
                if any(mtype.endswith(s) for s in VLLM_DELTA_SUFFIXES):
                    text = self._extract_text(msg)
                    if text:
                        await self._on_delta(sess, text)
                elif any(mtype.endswith(s) for s in VLLM_DONE_SUFFIXES):
                    text = self._extract_text(msg)
                    if text:
                        await self._on_delta(sess, text, final=True)
        except Exception:
            logger.debug("Stream %s: vLLM-Reader beendet", sess.request_id[:8])

    @staticmethod
    def _extract_text(msg: dict) -> str:
        for f in VLLM_DELTA_FIELDS:
            v = msg.get(f)
            if isinstance(v, str) and v:
                return v
        return ""

    async def _on_delta(self, sess: StreamSession, text: str, final: bool = False) -> None:
        """Neuer/finaler Transkript-Text vom vLLM. delta = inkrementell; wir
        haengen an, wenn er den bisherigen Partial verlaengert, sonst ersetzen
        wir (Voxtral kann korrigieren)."""
        if final or text.startswith(sess.last_partial):
            new_full = text if final else text
        else:
            new_full = (sess.last_partial + text).strip()
        new_full = new_full.strip()
        if new_full and new_full != sess.last_partial:
            sess.last_partial = new_full
            sess.last_growth_at = time.time()
            if self._ws is not None:
                await _send(self._ws, "stt_partial", {
                    "requestId": sess.request_id,
                    "audioRequestId": sess.audio_request_id,
                    "text": new_full,
                })

    def feed_chunk(self, payload: dict) -> bool:
        request_id = payload.get("requestId", "")
        sess = self._sessions.get(request_id)
        if sess is None or sess.closed:
            return False
        pcm_b64 = payload.get("pcm", "")
        if not pcm_b64:
            return True
        try:
            pcm = base64.b64decode(pcm_b64)
        except Exception:
            return True
        sess.pcm_buffer.extend(pcm)
        sess.last_chunk_at = time.time()
        # An vLLM weiterreichen (fire-and-forget).
        if sess.vllm_ws is not None:
            asyncio.create_task(self._vllm_append(sess, pcm_b64))
        return True

    async def _vllm_append(self, sess: StreamSession, pcm_b64: str) -> None:
        try:
            await sess.vllm_ws.send(json.dumps({
                "type": VLLM_SEND_APPEND,
                VLLM_AUDIO_FIELD: pcm_b64,
            }))
        except Exception:
            pass

    def end_session(self, request_id: str) -> None:
        sess = self._sessions.get(request_id)
        if sess is not None:
            sess.closed = True

    def drop(self, request_id: str) -> None:
        sess = self._sessions.pop(request_id, None)
        if sess is not None:
            self._teardown_vllm(sess)

    def _teardown_vllm(self, sess: StreamSession) -> None:
        try:
            if sess.vllm_reader is not None:
                sess.vllm_reader.cancel()
        except Exception:
            pass
        if sess.vllm_ws is not None:
            asyncio.create_task(self._close_ws(sess.vllm_ws))
            sess.vllm_ws = None

    @staticmethod
    async def _close_ws(ws) -> None:
        try:
            await ws.close()
        except Exception:
            pass

    # ── Endpointer (adaptiv, wie whisper-Bridge M0.1) ──
    def _buffer_duration_ms(self, sess: StreamSession) -> float:
        samples = len(sess.pcm_buffer) // 2
        return (samples / sess.sample_rate) * 1000.0 if samples else 0.0

    def _tail_rms(self, sess: StreamSession) -> float:
        win_bytes = int(sess.sample_rate * STREAM_ENERGY_WINDOW_MS / 1000) * 2
        if win_bytes <= 0:
            return 0.0
        tail = sess.pcm_buffer[-win_bytes:]
        if len(tail) < 2:
            return 0.0
        arr = pcm_s16le_to_float32(bytes(tail))
        if arr.size == 0:
            return 0.0
        return float(np.sqrt(np.mean(arr * arr)))

    def _voice_threshold(self, sess: StreamSession) -> float:
        nf = sess.noise_floor
        if nf <= 0.0:
            return sess.voice_rms_min
        return min(max(nf * sess.voice_factor, sess.voice_rms_min), sess.voice_rms_max)

    def _update_noise_floor(self, sess: StreamSession, rms: float) -> None:
        nf = sess.noise_floor
        if nf <= 0.0:
            sess.noise_floor = rms
        elif rms < nf:
            sess.noise_floor = 0.90 * nf + 0.10 * rms
        else:
            sess.noise_floor = 0.98 * nf + 0.02 * rms

    async def run_endpointer(self) -> None:
        logger.info("Voxtral-Endpointer gestartet (adaptiver VAD)")
        while True:
            await asyncio.sleep(0.2)
            now = time.time()
            for sid, sess in list(self._sessions.items()):
                try:
                    await self._tick(sess, now)
                except Exception:
                    logger.exception("Endpointer-Tick crashed (session=%s)", sid[:8])
            for sid, sess in list(self._sessions.items()):
                if now - sess.last_chunk_at > STREAM_SESSION_TTL_S:
                    logger.info("Stream %s: TTL — drop", sid[:8])
                    self.drop(sid)

    async def _tick(self, sess: StreamSession, now: float) -> None:
        if sess.endpoint_sent:
            return
        elapsed_ms = (now - sess.started_at) * 1000.0
        if elapsed_ms > sess.hard_cap_ms and not sess.closed:
            await self._finalize(sess, reason="hardcap")
            return
        if sess.closed:
            await self._finalize(sess, reason="stream_end")
            return
        if self._buffer_duration_ms(sess) < STREAM_MIN_AUDIO_MS:
            return
        # adaptive akustische Sprach-Aktivitaet
        rms = self._tail_rms(sess)
        if rms >= self._voice_threshold(sess):
            sess.last_voice_at = now
        else:
            self._update_noise_floor(sess, rms)
        # Endpoint: akustisch (Primaer) oder semantisch (Backstop), sobald Text da
        if sess.last_growth_at > 0.0:
            acoustic_silence_ms = (now - sess.last_voice_at) * 1000.0 if sess.last_voice_at > 0 else 0.0
            semantic_silence_ms = (now - sess.last_growth_at) * 1000.0
            acoustic_done = sess.last_voice_at > 0 and acoustic_silence_ms >= sess.endpoint_ms
            semantic_done = semantic_silence_ms >= sess.endpoint_ms * STREAM_SEMANTIC_BACKUP_FACTOR
            if acoustic_done or semantic_done:
                await self._finalize(sess, reason="endpoint" if acoustic_done else "endpoint_semantic")

    async def _finalize(self, sess: StreamSession, reason: str) -> None:
        if sess.endpoint_sent:
            return
        sess.endpoint_sent = True
        # vLLM ggf. committen, damit ein letztes transcription.done kommt.
        if sess.vllm_ws is not None:
            try:
                await sess.vllm_ws.send(json.dumps({"type": VLLM_SEND_COMMIT}))
                await asyncio.sleep(0.15)  # kurz auf finalen Delta warten
            except Exception:
                pass
        final_text = sess.last_partial.strip()
        duration_s = self._buffer_duration_ms(sess) / 1000.0
        logger.info("Stream %s: FINAL (reason=%s, %.1fs): %r",
                    sess.request_id[:8], reason, duration_s, final_text[:120])
        if self._ws is not None:
            endpoint_payload = {
                "requestId": sess.request_id,
                "audioRequestId": sess.audio_request_id,
                "text": final_text,
                "reason": reason,
                "durationS": duration_s,
                "sttMs": 0,
                "voice": sess.voice,
                "speed": sess.speed,
                "interrupted": sess.interrupted,
            }
            if sess.location:
                endpoint_payload["location"] = sess.location
            await _send(self._ws, "stt_endpoint", endpoint_payload)
            await _send(self._ws, "stt_stream_done", {
                "requestId": sess.request_id,
                "audioRequestId": sess.audio_request_id,
                "text": final_text,
                "reason": reason,
            })
        self.drop(sess.request_id)


async def _broadcast_status(ws, state: str, **extra) -> None:
    payload = {"service": "voxtral", "state": state}
    payload.update(extra)
    await _send(ws, "service_status", payload)


async def run_loop(sessions: SessionManager) -> None:
    use_tls = RVS_TLS
    retry_s = 2
    tls_fallback_tried = False
    while True:
        scheme = "wss" if use_tls else "ws"
        url = f"{scheme}://{RVS_HOST}:{RVS_PORT}/ws?token={RVS_TOKEN}"
        masked = url.replace(RVS_TOKEN, "***") if RVS_TOKEN else url
        try:
            logger.info("Verbinde zu RVS: %s", masked)
            async with websockets.connect(url, ping_interval=20, ping_timeout=10,
                                           max_size=50 * 1024 * 1024) as ws:
                logger.info("RVS verbunden")
                retry_s = 2
                tls_fallback_tried = False
                sessions.attach_ws(ws)
                await _broadcast_status(ws, "ready", model=VOXTRAL_MODEL)
                await _send(ws, "config_request", {"service": "voxtral"})

                async for raw in ws:
                    try:
                        msg = json.loads(raw)
                    except Exception:
                        continue
                    mtype = msg.get("type", "")
                    payload = msg.get("payload", {}) or {}
                    if mtype == "stt_stream_start":
                        asyncio.create_task(sessions.start_session(payload))
                    elif mtype == "stt_audio_chunk":
                        sessions.feed_chunk(payload)
                    elif mtype == "stt_stream_end":
                        sessions.end_session(payload.get("requestId", ""))
                    # stt_request (Legacy One-Shot) macht Voxtral hier NICHT —
                    # dafuer bleibt die whisper-Bridge (Fallback).
        except Exception as e:
            logger.warning("RVS-Verbindung verloren: %s — retry in %ds", e, retry_s)
            if use_tls and RVS_TLS_FALLBACK and not tls_fallback_tried:
                use_tls = False
                tls_fallback_tried = True
                logger.info("TLS-Fallback: versuche ws:// (kein TLS)")
                continue
            await asyncio.sleep(retry_s)
            retry_s = min(retry_s * 2, 30)
            use_tls = RVS_TLS  # fuer den naechsten Zyklus zuruecksetzen


async def main() -> None:
    if not RVS_HOST or not RVS_TOKEN:
        logger.error("RVS_HOST/RVS_TOKEN fehlen — .env pruefen. Abbruch.")
        return
    sessions = SessionManager()
    logger.info("Voxtral-Bridge startet — vLLM=%s Modell=%s", VOXTRAL_VLLM_URL, VOXTRAL_MODEL)
    await asyncio.gather(
        run_loop(sessions),
        sessions.run_endpointer(),
    )


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        pass