Bolt-Pet/bolt_pet/controller.py

"""Orchestrates the pet's mic -> wake word -> STT -> server -> TTS pipeline.

Runs on a background QThread (see ui/app.py) so the Qt event loop / window
painting is never blocked by audio I/O or network calls. Talks to the UI
only through Qt signals (state_changed / said / log / action / napping),
which Qt marshals safely across threads — this class never touches a QWidget
directly.

Beyond the core loop it owns the side channels that let the pet act on its
own: the heartbeat (proactive announcements), the desktop notification
bridge, quiet hours, barge-in, and the live wake-word threshold.
"""

from __future__ import annotations

import threading
import time
from typing import Optional

from PySide6.QtCore import QObject, Signal

from . import (
    config, history as history_mod, monitors as monitors_mod, notifications,
    pet_actions, quiet, screen_context, screen_text, server_client,
    speech_text, updater,
)
from .audio import barge_in, mic, stt, tts, wake_word
from .state import PetState, PetStateMachine

# How often to re-check whether the pet should be napping. The fullscreen
# probe shells out to xprop, so this deliberately isn't every heartbeat tick.
_NAP_CHECK_INTERVAL_SECONDS = 10.0


class PetController(QObject):
    state_changed = Signal(str)  # PetState.value
    said = Signal(str)  # text now showing in the speech bubble
    log = Signal(str)
    action = Signal(dict)  # parsed petctl action for the UI to perform
    napping = Signal(bool)  # quiet hours / fullscreen do-not-disturb
    restart_requested = Signal(str)  # version we just updated to
    finished = Signal()

    def __init__(self):
        super().__init__()
        self._running = True
        self._muted = False
        self._talk_now = threading.Event()
        self._stream = None
        self._last_heartbeat = 0.0
        self._state = PetStateMachine(on_change=self._handle_state_change)

        # Conversation scrollback, shared read-only with the UI's History
        # window. Append-only from this thread; the UI only ever snapshots it.
        self.history = history_mod.ConversationHistory(limit=config.HISTORY_LIMIT)

        # The screen layout, as published by the UI (see set_monitors). Held
        # here rather than probed, so "monitor 2" means the same thing to the
        # controller and to the window that has to jump there — see
        # monitors.py for why that matters.
        self._monitors: list[monitors_mod.Monitor] = []
        self._pet_monitor: Optional[int] = None

        # Wake-word sensitivity is live-tunable (tray tuner), so it's read
        # through a callable on every frame rather than captured per listen.
        self._wake_threshold = config.WAKE_WORD_THRESHOLD
        self._near_misses = wake_word.NearMissLog()

        self._barge_in: Optional[barge_in.BargeInDetector] = None
        self._napping = False
        self._nap_forced: Optional[bool] = None  # petctl nap on/off overrides the schedule
        self._last_nap_check = 0.0

        # When a reply ends on a question the pet keeps listening for the
        # answer. _follow_ups counts how many have chained without you
        # re-triggering, so a server that ends every reply with "?" can't
        # loop forever off mic noise.
        self._pending_follow_up = False
        self._follow_ups = 0

        self._last_update_check = 0.0
        self._update_pending = False  # applied on disk, waiting for the restart

        self._notification_watcher: Optional[notifications.NotificationWatcher] = None
        self._notification_gate = notifications.NotificationGate(
            config.NOTIFICATION_FILTER, config.NOTIFICATION_MIN_INTERVAL_SECONDS
        )
        self._pending_notifications: list[notifications.Notification] = []
        self._notification_lock = threading.Lock()

    # ── external controls (safe to call from the Qt/UI thread) ─────────

    def request_talk_now(self) -> None:
        self._talk_now.set()

    def toggle_mute(self) -> bool:
        self._muted = not self._muted
        self.log.emit("Muted." if self._muted else "Unmuted.")
        return self._muted

    def set_napping(self, napping: Optional[bool]) -> None:
        """Force the nap state on/off, or pass None to hand control back to
        the quiet-hours schedule."""
        self._nap_forced = napping
        if napping is not None:
            self._apply_nap_state(napping)

    def wake_threshold(self) -> float:
        return self._wake_threshold

    def set_wake_threshold(self, value: float) -> None:
        self._wake_threshold = min(max(float(value), 0.01), 0.99)

    def wake_stats(self) -> dict:
        return {"peak": self._near_misses.peak, "near_misses": self._near_misses.entries()}

    def reset_wake_stats(self) -> None:
        self._near_misses.clear()

    def stop(self) -> None:
        self._running = False
        self._talk_now.set()  # wake up anything blocked waiting on it
        if self._notification_watcher is not None:
            self._notification_watcher.stop()

    # ── internal ─────────────────────────────────────────────────────────

    def _handle_state_change(self, _old: PetState, new: PetState) -> None:
        self.state_changed.emit(new.value)

    def _should_continue(self) -> bool:
        return self._running

    def run(self) -> None:
        """Thread entry point (connected to QThread.started)."""
        missing = config.missing_config()
        if missing:
            self.log.emit(f"Missing config: {', '.join(missing)} — set them in .env and restart.")
            self.finished.emit()
            return
        try:
            self._stream = mic.open_input_stream()
        except Exception as exc:
            self.log.emit(f"Could not open microphone: {exc}")
            self.finished.emit()
            return

        if config.BARGE_IN:
            # In wake mode the detector shares the idle listener's model and
            # its live threshold, so the tray tuner's slider applies to
            # interrupting as well as waking (unless BARGE_IN_WAKE_THRESHOLD
            # pins it to a fixed, stricter number).
            self._barge_in = barge_in.make_detector(
                self._stream,
                wake_threshold=(
                    config.BARGE_IN_WAKE_THRESHOLD
                    if config.BARGE_IN_WAKE_THRESHOLD is not None
                    else self.wake_threshold
                ),
            )
            self.log.emit(f"Barge-in: {config.BARGE_IN_MODE} mode.")

        with self._stream:
            try:
                health = server_client.check_health()
                self.log.emit(f"Connected to server: {health}")
            except Exception as exc:
                self.log.emit(f"Server not reachable yet ({exc}) — will keep trying per-request.")
            self._start_notification_bridge()
            self._loop()
        if self._notification_watcher is not None:
            self._notification_watcher.stop()
        self.finished.emit()

    def _loop(self) -> None:
        while self._running:
            if self._muted:
                triggered = self._talk_now.wait(timeout=0.5)
                if not self._running:
                    return
                if not triggered:
                    continue
                self._talk_now.clear()
            else:
                if not self._wait_for_wake_or_click():
                    if not self._running:
                        return
                    continue
            self._handle_conversation_turn()

    def _wait_for_wake_or_click(self) -> bool:
        """True once either the wake phrase was heard or a click-to-talk
        request came in; False on a spurious wakeup (loop again)."""
        def should_continue() -> bool:
            return self._running and not self._talk_now.is_set()

        detected = wake_word.listen_for_wake_word(
            self._stream,
            should_continue=should_continue,
            threshold=self.wake_threshold,  # callable: the tuner slider is live
            on_tick=self._maybe_heartbeat,
            on_score=self._observe_wake_score,
        )
        if not self._running:
            return False
        if detected:
            self._talk_now.clear()  # in case both fired around the same time
            return True
        if self._talk_now.is_set():
            self._talk_now.clear()
            return True
        return False

    def _observe_wake_score(self, score: float, threshold: float) -> None:
        self._near_misses.observe(score, threshold, time.time())

    def _handle_conversation_turn(self) -> None:
        # A turn you started yourself ends any follow-up chain in progress.
        following_up, self._pending_follow_up = self._pending_follow_up, False
        if not following_up:
            self._follow_ups = 0

        self._state.transition(PetState.LISTENING)
        pcm = mic.record_utterance(
            self._stream,
            should_continue=self._should_continue,
            # Answering a question deserves longer than saying the wake word
            # on purpose does — you were just asked something.
            grace_s=config.FOLLOW_UP_GRACE_SECONDS if following_up else None,
        )
        if pcm is None:
            self._follow_ups = 0  # silence ends the chain
            self._state.transition(PetState.IDLE)
            return

        self._state.transition(PetState.THINKING)
        try:
            text = stt.transcribe(pcm)
        except stt.SttError as exc:
            self.log.emit(f"STT failed: {exc}")
            self._state.transition(PetState.ERROR)
            self._state.transition(PetState.IDLE)
            return
        if not text:
            self._state.transition(PetState.IDLE)
            return
        self.log.emit(f"You: {text}")
        self.history.add(history_mod.USER, text, time.time())

        try:
            # What's focused right now rides along, so "what's this error?"
            # has a referent without you having to describe the window.
            reply = server_client.converse(
                self._with_context(text), on_command=self._handle_command
            )
        except server_client.ServerError as exc:
            self.log.emit(f"Server error: {exc}")
            self._state.transition(PetState.ERROR)
            self._state.transition(PetState.IDLE)
            return

        self._speak(reply)
        self._state.transition(PetState.IDLE)

    def _with_context(self, text: str) -> str:
        """Everything the server gets alongside what you actually said: the
        focused window title, and a one-line note about the screen layout so
        Bolt knows how many monitors there are and where he's standing
        without having to ask. Only the *layout* rides along for free — the
        text on those screens costs an OCR pass, so it stays behind
        `petctl read`."""
        text = screen_context.context_for(text)
        if config.MONITOR_CONTEXT:
            text = monitors_mod.annotate(text, self._monitors, self._pet_monitor)
        return text

    # ── screen layout, published by the UI ───────────────────────────────

    def set_monitors(self, monitors: list) -> None:
        """Slot: the window telling us what screens exist (queued signal)."""
        self._monitors = list(monitors)
        self.log.emit(
            "Screens: " + (monitors_mod.summary(self._monitors) or "none reported")
        )

    def set_pet_monitor(self, index: int) -> None:
        """Slot: the window telling us which screen the pet is standing on."""
        self._pet_monitor = int(index)

    def _handle_command(self, command: str) -> str:
        """Server-relayed command. `petctl ...` drives the pet's body and
        never reaches a shell; everything else is a real command, exactly as
        before (see the security notes in the README)."""
        try:
            action = pet_actions.parse(command)
        except pet_actions.ActionError as exc:
            self.log.emit(f"petctl: {exc}")
            return f"[pet] {exc}"
        if action is None:
            return server_client.run_local_command(command)
        self.log.emit(f"Pet action: {action}")

        # Queries answer from here rather than from pet_actions.describe():
        # their output *is* the useful part, and it's what the server reads
        # back off the tool-result relay.
        kind = action["action"]
        if kind == "monitors":
            return monitors_mod.describe(self._monitors, self._pet_monitor)
        if kind == "read":
            return self._read_screen(action["target"])
        if kind == "jump":
            try:
                target = monitors_mod.resolve(
                    self._monitors, action["target"], self._pet_monitor
                )
            except ValueError as exc:
                self.log.emit(f"petctl jump: {exc}")
                return f"[pet] {exc}"
            # Hand the window a resolved index, so it can't re-resolve the
            # spec against a different screen ordering.
            self.action.emit({"action": "jump", "monitor": target.index})
            return f"[pet] jumped to monitor {target.label}"

        if kind == "nap":
            self.set_napping(bool(action["enabled"]))
        self.action.emit(action)
        return pet_actions.describe(action)

    def _read_screen(self, target: str) -> str:
        """`petctl read` — OCR a screen and hand the text back to the server."""
        if not config.SCREEN_TEXT:
            return "[pet] screen reading is disabled (set SCREEN_TEXT=true in .env)"
        if not self._monitors:
            return "[pet] no monitor information available"
        limit = config.SCREEN_TEXT_MAX_CHARS
        if target in ("all", "everything", "*"):
            self.log.emit(f"Reading all {len(self._monitors)} screens…")
            return screen_text.read_monitors(self._monitors, limit)
        if target in ("here", "", "this", "current"):
            index = self._pet_monitor if self._pet_monitor is not None else 0
            monitor = self._monitors[min(index, len(self._monitors) - 1)]
        else:
            try:
                monitor = monitors_mod.resolve(
                    self._monitors, target, self._pet_monitor
                )
            except ValueError as exc:
                return f"[pet] {exc}"
        self.log.emit(f"Reading monitor {monitor.number} ({monitor.name})…")
        return screen_text.read_monitor(monitor, limit)

    def _speak(self, text: str) -> None:
        self._state.transition(PetState.TALKING)
        # Bubble gets the markdown stripped but emoji kept (it can't render
        # **bold** but draws emoji fine); tts.speak() does its own, stricter
        # sanitizing for the voice.
        self.said.emit(speech_text.for_display(text))
        self.log.emit(f"Bolt: {text}")
        self.history.add(history_mod.PET, text, time.time())

        should_stop = None
        if self._barge_in is not None:
            self._barge_in.reset()
            should_stop = self._barge_in.check
        completed = tts.speak(
            text,
            on_error=lambda exc: self.log.emit(f"TTS failed: {exc}"),
            should_stop=should_stop,
        )
        # Read the scoring history *before* resetting, or the log reports the
        # blank counters instead of what actually fired.
        detail = self._barge_in_detail()
        if self._barge_in is not None:
            # Playback fed the pet's own voice into the wake model's rolling
            # window. Clear it before the idle listener starts scoring again,
            # or Bolt's last sentence is still in there being re-scored.
            self._barge_in.reset()
        if not completed:
            # You talked over it — take that as the start of the next turn
            # rather than making you say the wake word again.
            self.log.emit(f"Interrupted — listening. {detail}")
            self._follow_ups = 0  # you're clearly engaged; start the count over
            self._talk_now.set()
        elif self._should_follow_up(text):
            self._follow_ups += 1
            cap = config.FOLLOW_UP_MAX_TURNS
            self.log.emit(
                f"Asked a question — listening for your answer "
                f"({self._follow_ups}{'/' + str(cap) if cap > 0 else ''})."
            )
            self._pending_follow_up = True
            self._talk_now.set()

    def _should_follow_up(self, text: str) -> bool:
        """Whether *text* leaves the pet waiting on an answer.

        Muted is excluded because mute means "don't listen to me" — an
        automatic turn would walk straight past it. Napping isn't: quiet
        hours suppress the pet *starting* something, and a question is only
        ever asked in reply to you."""
        if not config.FOLLOW_UP_LISTEN or self._muted:
            return False
        if not speech_text.is_question(text):
            return False
        # Only worth mentioning the cap on a reply that would otherwise have
        # kept listening, or it fires on every statement the pet makes.
        if config.FOLLOW_UP_MAX_TURNS > 0 and self._follow_ups >= config.FOLLOW_UP_MAX_TURNS:
            self.log.emit("Follow-up limit reached — say the wake word to keep going.")
            return False
        return True

    def _barge_in_detail(self) -> str:
        """Why the interruption fired, for the log. How far into playback it
        happened is the tell: frame 1 means the detector was still holding
        audio from before this reply started, whereas a hit several seconds
        in is something the mic actually heard."""
        detector = self._barge_in
        if isinstance(detector, barge_in.WakeWordBargeIn):
            return (
                f"(wake score {detector.last_score:.3f} >= {detector.last_threshold:.2f}, "
                f"peak {detector.peak_score:.3f}, at frame {detector.frames_checked} / "
                f"{detector.seconds_checked:.1f}s into playback)"
            )
        if isinstance(detector, barge_in.BargeInDetector):
            return f"(loud frames {detector.loud_frames}, threshold {config.BARGE_IN_RMS_THRESHOLD})"
        return ""

    # ── quiet hours / do-not-disturb ─────────────────────────────────────

    def _apply_nap_state(self, napping: bool) -> None:
        if napping == self._napping:
            return
        self._napping = napping
        self.log.emit("Napping — no proactive noise." if napping else "Awake.")
        self.napping.emit(napping)

    def _refresh_nap_state(self) -> None:
        if self._nap_forced is not None:
            self._apply_nap_state(self._nap_forced)
            return
        now = time.monotonic()
        if now - self._last_nap_check < _NAP_CHECK_INTERVAL_SECONDS:
            return
        self._last_nap_check = now
        napping = quiet.is_quiet(
            config.QUIET_HOURS,
            on_error=lambda exc: self.log.emit(f"QUIET_HOURS is malformed ({exc}) — ignoring it."),
        )
        if not napping and config.DND_ON_FULLSCREEN:
            napping = screen_context.is_fullscreen_active()
        self._apply_nap_state(napping)

    # ── desktop notification bridge ──────────────────────────────────────

    def _start_notification_bridge(self) -> None:
        if not config.NOTIFICATION_BRIDGE:
            return
        watcher = notifications.NotificationWatcher(self._queue_notification)
        problem = watcher.start()
        if problem:
            self.log.emit(problem)
            return
        self._notification_watcher = watcher
        self.log.emit("Notification bridge on.")

    def _queue_notification(self, notification: notifications.Notification) -> None:
        """Called on the watcher thread — just queue it; forwarding happens on
        the pipeline thread where it can't collide with a live conversation."""
        if not self._notification_gate.should_forward(notification, time.monotonic()):
            return
        with self._notification_lock:
            self._pending_notifications.append(notification)

    def _drain_notifications(self) -> None:
        with self._notification_lock:
            pending, self._pending_notifications = self._pending_notifications, []
        for notification in pending:
            if not self._running or self._napping:
                return
            self.log.emit(f"Notification: {notification.as_text()}")
            self.history.add(history_mod.SYSTEM, notification.as_text(), time.time())
            try:
                reply = server_client.converse(
                    f"[desktop notification] {notification.as_text()}",
                    on_command=self._handle_command,
                )
            except server_client.ServerError as exc:
                self.log.emit(f"Couldn't forward notification: {exc}")
                return
            if reply.strip():
                self._speak(reply)
                self._state.transition(PetState.IDLE)

    # ── auto-update ──────────────────────────────────────────────────────

    def _maybe_update(self) -> None:
        """Poll the Gitea releases page and, if there's a newer tag, apply it
        and ask the UI to restart.

        Only ever runs from the wake-listener's tick, so the pet is IDLE and
        between turns by construction — an update can't land mid-sentence.
        Failures are logged and the interval resets, so a server that's down
        (or a release that rolls back) costs one log line an hour, not a
        retry storm."""
        if not config.AUTO_UPDATE or self._update_pending:
            return
        now = time.monotonic()
        if now - self._last_update_check < config.UPDATE_CHECK_INTERVAL_SECONDS:
            return
        self._last_update_check = now
        try:
            release = updater.check_for_update()
        except updater.UpdateError as exc:
            self.log.emit(f"Update check failed: {exc}")
            return
        if release is None:
            return

        self.log.emit(f"Update available: {release.tag} — applying.")
        try:
            previous = updater.apply_update(release.tag, on_log=self.log.emit)
        except updater.UpdateError as exc:
            self.log.emit(f"Update to {release.tag} failed: {exc}")
            return

        self._update_pending = True
        self.log.emit(f"Updated {previous} -> {release.tag}; restarting.")
        if not self._napping:
            # Napping means no proactive noise, so a silent restart it is.
            self._speak(f"Updating to {release.tag}. Back in a second.")
            self._state.transition(PetState.IDLE)
        self.restart_requested.emit(release.tag)

    # ── heartbeat ────────────────────────────────────────────────────────

    def _maybe_heartbeat(self) -> None:
        self._refresh_nap_state()
        self._maybe_update()
        if self._update_pending:
            return  # on the way out — don't start a conversation now
        now = time.monotonic()
        if now - self._last_heartbeat < config.HEARTBEAT_INTERVAL_SECONDS:
            return
        self._last_heartbeat = now
        if self._state.state != PetState.IDLE:
            return
        if self._napping:
            return  # quiet hours: still answers when spoken to, just doesn't start
        self._drain_notifications()
        if self._state.state != PetState.IDLE:
            return
        try:
            announcement = server_client.report_status()
        except server_client.ServerError as exc:
            self.log.emit(f"Heartbeat failed: {exc}")
            return
        if announcement:
            self._speak(announcement)
            self._state.transition(PetState.IDLE)