STT-Assistant/src/stt_daemon/daemon.py

"""Main daemon that orchestrates capture, segmentation, and transcription."""

import logging
import queue
import threading
from typing import Optional

from .audio import AudioCapture, open_capture
from .config import Config
from .output import TranscriptSink
from .segmenter import FrameSegmenter, make_vad
from .transcriber import Transcriber


class Daemon:
    """Always-on STT daemon with worker threads."""

    def __init__(self, cfg: Config, log_level: int = logging.INFO):
        """Initialize the daemon.

        Args:
            cfg: Configuration object.
            log_level: Logging level (e.g., logging.INFO, logging.DEBUG).
        """
        self.cfg = cfg

        # Set up logging
        self.logger = logging.getLogger("stt_daemon")
        self.logger.setLevel(log_level)
        if not self.logger.handlers:
            handler = logging.StreamHandler()
            handler.setFormatter(logging.Formatter("[%(levelname)s] %(message)s"))
            self.logger.addHandler(handler)

        # Initialize components
        self.capture: Optional[AudioCapture] = None
        self.segmenter: Optional[FrameSegmenter] = None
        self.transcriber: Optional[Transcriber] = None
        self.sink: Optional[TranscriptSink] = None

        # Threading
        self.segment_queue: queue.Queue = queue.Queue(maxsize=50)
        self._stop_event = threading.Event()
        self._stopped = False
        self._reader_thread: Optional[threading.Thread] = None
        self._worker_thread: Optional[threading.Thread] = None

    def run(self) -> None:
        """Run the daemon (blocks until stopped)."""
        try:
            self._initialize()
            self._start_workers()
            self._stop_event.wait()
        finally:
            self.stop()

    def _initialize(self) -> None:
        """Initialize all components."""
        self.logger.info("Initializing audio capture...")
        self.capture = open_capture(self.cfg.audio)

        self.logger.info("Initializing VAD...")
        vad = make_vad(self.cfg.vad.aggressiveness)

        self.logger.info("Initializing segmenter...")
        self.segmenter = FrameSegmenter(
            sample_rate=self.cfg.audio.sample_rate,
            frame_ms=self.cfg.audio.frame_ms,
            vad=vad,
            pre_roll_ms=self.cfg.vad.pre_roll_ms,
            speech_start_ms=self.cfg.vad.speech_start_ms,
            silence_ms=self.cfg.vad.silence_ms,
            min_utterance_ms=self.cfg.vad.min_utterance_ms,
            max_utterance_s=self.cfg.vad.max_utterance_s,
        )

        self.logger.info(f"Loading Whisper model '{self.cfg.whisper.model}'...")
        self.transcriber = Transcriber(
            model=self.cfg.whisper.model,
            device=self.cfg.whisper.device,
            compute_type=self.cfg.whisper.compute_type,
            language=self.cfg.whisper.language,
            beam_size=self.cfg.whisper.beam_size,
            vad_filter=self.cfg.whisper.vad_filter,
        )

        self.logger.info("Initializing output sink...")
        self.sink = TranscriptSink(
            stdout=self.cfg.output.stdout,
            jsonl_dir=self.cfg.output.jsonl_dir,
        )

        self.logger.info("Daemon ready. Listening...")

    def _start_workers(self) -> None:
        """Start reader and worker threads."""
        self._reader_thread = threading.Thread(target=self._reader_loop, daemon=False)
        self._worker_thread = threading.Thread(target=self._worker_loop, daemon=False)
        self._reader_thread.start()
        self._worker_thread.start()

    def _reader_loop(self) -> None:
        """Reader thread: capture frames and segment them."""
        try:
            if not self.capture or not self.segmenter:
                return

            for frame in self.capture.frames():
                if self._stop_event.is_set():
                    break

                utt = self.segmenter.push(frame)
                if utt:
                    try:
                        self.segment_queue.put(utt, block=False)
                    except queue.Full:
                        self.logger.warning("Segment queue full, dropping oldest")
                        try:
                            self.segment_queue.get_nowait()
                            self.segment_queue.put(utt, block=False)
                        except queue.Empty:
                            pass

            # Shutdown: flush any in-progress utterance
            if self.segmenter:
                utt = self.segmenter.flush()
                if utt:
                    try:
                        self.segment_queue.put(utt, block=False)
                    except queue.Full:
                        pass

        except Exception as e:
            if self._stop_event.is_set():
                self.logger.debug(f"Reader stopped: {e}")
            else:
                self.logger.error(f"Reader thread error: {e}")
        finally:
            if self.capture:
                self.capture.close()

    def _worker_loop(self) -> None:
        """Worker thread: transcribe segments and emit results."""
        try:
            while not self._stop_event.is_set():
                try:
                    utt = self.segment_queue.get(timeout=1)
                except queue.Empty:
                    continue

                if self.transcriber and self.sink:
                    try:
                        result = self.transcriber.transcribe(utt.audio)
                        self.sink.emit(result, utt.duration_s)
                    except Exception as e:
                        self.logger.error(f"Transcription error: {e}")

        except Exception as e:
            self.logger.error(f"Worker thread error: {e}")

    def stop(self) -> None:
        """Stop the daemon. Safe to call more than once."""
        if self._stopped:
            return
        self._stopped = True

        self.logger.info("Stopping daemon...")
        self._stop_event.set()

        # Close capture FIRST so the reader thread's blocking read unblocks
        # (otherwise a non-daemon reader thread can hang process shutdown).
        if self.capture:
            self.capture.close()

        if self._reader_thread:
            self._reader_thread.join(timeout=5)
        if self._worker_thread:
            self._worker_thread.join(timeout=5)

        if self.transcriber:
            self.transcriber.close()
        if self.sink:
            self.sink.close()

        self.logger.info("Daemon stopped.")