Deaf_Transcription_Service/bridge/bridge.py

#!/usr/bin/env python3
"""
bridge.py — Live Transcription Bridge

Uses WhisperLiveKit's AudioProcessor API directly (no WebSocket).
Publishes rolling 3-line JSON to Mosquitto MQTT.

Run this script:
    python bridge.py
"""

import asyncio
import json
import queue as _stdlib_queue
import re
import textwrap
import threading
import time
from collections import Counter
from pathlib import Path

import numpy as np
import paho.mqtt.client as mqtt
import sounddevice as sd

from fastapi import FastAPI, Request
from whisperlivekit import AudioProcessor, TranscriptionEngine
import uvicorn

# ── Configuration ─────────────────────────────────────────────────────────────

MQTT_HOST        = "localhost"
MQTT_PORT        = 1883
MQTT_TOPIC_TEXT  = "display/text"
MQTT_TOPIC_CLEAR = "display/clear"

SAMPLE_RATE = 16000
CHANNELS    = 1
BLOCKSIZE   = 4096          # ~256 ms per chunk at 16 kHz

SENTENCE_TIMEOUT = 4.0      # seconds of silence before forcing a flush
MAX_LINE_CHARS   = 38       # characters per line
DISPLAY_LINES    = 3

# Set to a device index (integer) to force a specific microphone.
# Leave as None to use the Windows default input device.
AUDIO_DEVICE: int | None = 12

SPEAKERS_FILE = Path(__file__).parent / "speakers.json"

DEFAULT_SPEAKERS: dict[str, str] = {
    "SPEAKER_00": "Pastor",
    "SPEAKER_01": "Reader",
    "SPEAKER_02": "Guest",
    "SPEAKER_03": "Choir",
}

# ── Audio injection queue ─────────────────────────────────────────────────────
# stdlib queue.Queue is thread-safe across event loops; asyncio.Queue is not.
# admin.py POSTs test audio chunks to /inject (port 8002) which puts them here.
# _send_audio() drains this queue in preference to the live microphone.
_inject_queue: _stdlib_queue.Queue = _stdlib_queue.Queue(maxsize=240)

# ── Audio injection API ───────────────────────────────────────────────────────

_bridge_app = FastAPI()

@_bridge_app.post("/inject")
async def inject_audio(request: Request):
    chunk = await request.body()
    if chunk:
        try:
            _inject_queue.put_nowait(chunk)
        except _stdlib_queue.Full:
            pass
    return {"ok": True}

@_bridge_app.post("/inject/clear")
async def inject_clear():
    while True:
        try:
            _inject_queue.get_nowait()
        except _stdlib_queue.Empty:
            break
    return {"ok": True}


# ── Speaker persistence ───────────────────────────────────────────────────────

def _load_speakers() -> dict[str, str]:
    if SPEAKERS_FILE.exists():
        try:
            data = json.loads(SPEAKERS_FILE.read_text(encoding="utf-8"))
            if isinstance(data, dict):
                return data
        except (json.JSONDecodeError, OSError):
            pass
    _write_speakers(DEFAULT_SPEAKERS)
    return dict(DEFAULT_SPEAKERS)


def _write_speakers(names: dict[str, str]) -> None:
    try:
        SPEAKERS_FILE.write_text(
            json.dumps(names, indent=2, ensure_ascii=False),
            encoding="utf-8",
        )
    except OSError as exc:
        print(f"[Speakers] Save failed: {exc}")


# ── State ─────────────────────────────────────────────────────────────────────

class BridgeState:
    """All mutable state, protected by a single lock."""

    def __init__(self):
        self._lock                         = threading.Lock()
        self.speaker_names: dict[str, str] = _load_speakers()
        self._seen: set[str]               = set(self.speaker_names)
        self._current_speaker: str | None  = None
        self._speaker_changed              = False
        self._text_buffer                  = ""
        self._display: list[str]           = [""] * DISPLAY_LINES
        self._last_final_time              = time.monotonic()

    def set_speaker_name(self, speaker_id: str, name: str) -> None:
        with self._lock:
            self.speaker_names[speaker_id] = name.strip()
            self._seen.add(speaker_id)
            _write_speakers(self.speaker_names)

    def delete_speaker(self, speaker_id: str) -> None:
        with self._lock:
            self.speaker_names.pop(speaker_id, None)
            self._seen.discard(speaker_id)
            _write_speakers(self.speaker_names)

    def seen_speakers_snapshot(self) -> set[str]:
        with self._lock:
            return set(self._seen)

    def _resolve(self, speaker_id: str | None) -> str | None:
        if not speaker_id:
            return None
        return self.speaker_names.get(speaker_id, speaker_id)

    def push_final(self, text: str, speaker_id: str | None, mqtt_client: mqtt.Client) -> None:
        with self._lock:
            if speaker_id:
                self._seen.add(speaker_id)

            resolved = self._resolve(speaker_id)
            if resolved != self._current_speaker:
                if self._text_buffer:
                    self._flush(mqtt_client)
                self._current_speaker = resolved
                self._speaker_changed = True

            sep = " " if self._text_buffer else ""
            self._text_buffer += sep + text.strip()
            self._last_final_time = time.monotonic()

            if _is_sentence_end(text):
                self._flush(mqtt_client)

    def maybe_timeout_flush(self, mqtt_client: mqtt.Client) -> None:
        with self._lock:
            if self._text_buffer and (time.monotonic() - self._last_final_time) > SENTENCE_TIMEOUT:
                self._flush(mqtt_client)

    def _flush(self, mqtt_client: mqtt.Client) -> None:
        text = self._text_buffer.strip()
        self._text_buffer = ""
        if not text:
            return

        new_lines: list[str] = []
        if self._speaker_changed and self._current_speaker:
            new_lines.append(f"[{self._current_speaker.upper()}]")
            self._speaker_changed = False

        new_lines.extend(textwrap.wrap(text, MAX_LINE_CHARS) or [""])
        self._display.extend(new_lines)
        self._display = self._display[-DISPLAY_LINES:]
        while len(self._display) < DISPLAY_LINES:
            self._display.insert(0, "")

        payload = json.dumps({"lines": list(self._display)})
        mqtt_client.publish(MQTT_TOPIC_TEXT, payload)
        print(f"[Display] {self._display}")

    def clear(self, mqtt_client: mqtt.Client) -> None:
        with self._lock:
            self._display         = [""] * DISPLAY_LINES
            self._text_buffer     = ""
            self._current_speaker = None
            self._speaker_changed = False
        mqtt_client.publish(MQTT_TOPIC_CLEAR, "")
        print("[Display] Cleared")


# ── Helpers ───────────────────────────────────────────────────────────────────

def _is_sentence_end(text: str) -> bool:
    return bool(re.search(r'[.!?…]\s*$', text.strip()))


# ── MQTT ──────────────────────────────────────────────────────────────────────

def build_mqtt_client() -> mqtt.Client:
    client = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2)

    def on_connect(client, userdata, flags, rc, props):
        print("[MQTT] Connected" if rc == 0 else f"[MQTT] Failed: {rc}")

    def on_disconnect(client, userdata, flags, rc, props):
        print(f"[MQTT] Disconnected ({rc}), will reconnect...")

    client.on_connect    = on_connect
    client.on_disconnect = on_disconnect
    client.reconnect_delay_set(min_delay=1, max_delay=30)
    client.connect_async(MQTT_HOST, MQTT_PORT)
    client.loop_start()
    return client


# ── Audio pipeline ────────────────────────────────────────────────────────────

def _choose_audio_device() -> int | None:
    try:
        devices    = sd.query_devices()
        default_in = sd.default.device[0]
    except Exception as exc:
        print(f"[Audio] Cannot query devices: {exc}")
        return None

    print("[Audio] Available input devices:")
    input_devices: list[tuple[int, str]] = []
    for i, dev in enumerate(devices):
        if dev["max_input_channels"] > 0:
            marker = "  ← default" if i == default_in else ""
            print(f"  [{i}] {dev['name']}{marker}")
            input_devices.append((i, dev["name"]))

    if not input_devices:
        print("[Audio] ERROR: No input devices found.")
        return None

    if AUDIO_DEVICE is not None:
        print(f"[Audio] Using configured device [{AUDIO_DEVICE}]")
        return AUDIO_DEVICE

    if default_in >= 0:
        print(f"[Audio] Using default input device [{default_in}]")
        return default_in

    idx, name = input_devices[0]
    print(f"[Audio] No system default — using [{idx}] {name}")
    return idx


async def audio_processor_loop(state: BridgeState, mqtt_client: mqtt.Client, engine: TranscriptionEngine) -> None:
    audio_queue: asyncio.Queue[bytes] = asyncio.Queue(maxsize=120)
    loop = asyncio.get_running_loop()

    def audio_callback(indata: np.ndarray, frames: int, time_info, status) -> None:
        if status:
            print(f"[Audio] {status}")
        chunk = indata.tobytes()
        loop.call_soon_threadsafe(
            lambda: audio_queue.put_nowait(chunk) if not audio_queue.full() else None
        )

    device = _choose_audio_device()
    if device is None:
        print("[Audio] No input device — cannot start.")
        return

    audio_processor    = AudioProcessor(transcription_engine=engine)
    results_generator  = await audio_processor.create_tasks()

    async def _receive_results():
        # FrontData.lines is a cumulative list of committed Segment objects.
        # Track how many we've already processed so we only push new ones.
        seen_lines = 0
        async for response in results_generator:
            lines = response.lines or []
            # Guard against unexpected shrink (e.g. processor reset)
            if len(lines) < seen_lines:
                seen_lines = 0
            for seg in lines[seen_lines:]:
                text = (seg.text or "").strip()
                if text and not seg.is_silence():
                    spk = seg.speaker
                    speaker_id = f"SPEAKER_{spk:02d}" if isinstance(spk, int) and spk >= 0 else None
                    print(f"[Whisper] ({speaker_id or '?'}) {text}")
                    state.push_final(text, speaker_id, mqtt_client)
            seen_lines = len(lines)

    async def _send_audio():
        with sd.InputStream(
            device=device, samplerate=SAMPLE_RATE, channels=CHANNELS,
            dtype="int16", blocksize=BLOCKSIZE, callback=audio_callback,
        ):
            while True:
                # Injected test audio takes priority over live microphone
                try:
                    chunk = _inject_queue.get_nowait()
                except _stdlib_queue.Empty:
                    chunk = await audio_queue.get()
                await audio_processor.process_audio(chunk)

    flusher  = asyncio.create_task(_flusher(state, mqtt_client))
    reloader = asyncio.create_task(_speaker_reloader(state))
    try:
        await asyncio.gather(_send_audio(), _receive_results())
    finally:
        flusher.cancel()
        reloader.cancel()


async def _flusher(state: BridgeState, mqtt_client: mqtt.Client) -> None:
    while True:
        await asyncio.sleep(1.0)
        state.maybe_timeout_flush(mqtt_client)


async def _speaker_reloader(state: BridgeState) -> None:
    last_mtime = 0.0
    while True:
        await asyncio.sleep(5.0)
        try:
            mtime = SPEAKERS_FILE.stat().st_mtime
            if mtime != last_mtime:
                fresh = _load_speakers()
                with state._lock:
                    state.speaker_names = fresh
                last_mtime = mtime
                print("[Bridge] Speaker names reloaded from disk")
        except OSError:
            pass


# ── Entry point ───────────────────────────────────────────────────────────────

def main() -> None:
    state       = BridgeState()
    mqtt_client = build_mqtt_client()
    engine      = TranscriptionEngine(model_size="large-v3", lan="en", diarization=False, pcm_input=True)

    # Inject API must start before the audio loop so test playback works immediately
    def _run_inject_api():
        uvicorn.run(_bridge_app, host="127.0.0.1", port=8002, log_level="warning")

    inject_thread = threading.Thread(target=_run_inject_api, daemon=True)
    inject_thread.start()
    print("[Bridge] Test audio inject API at http://127.0.0.1:8002")

    def _run():
        asyncio.run(audio_processor_loop(state, mqtt_client, engine))

    ws_thread = threading.Thread(target=_run, daemon=True)
    ws_thread.start()
    print(f"[Bridge] Speaker names loaded from {SPEAKERS_FILE}")
    print("[Bridge] Audio pipeline running — speaker admin at http://localhost:8001")
    print("[Bridge] Close this window to quit")

    try:
        ws_thread.join()
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()