STT-Assistant/src/stt_daemon/tui/app.py

"""Textual TUI: watch transcriptions and drive the assistant."""

import argparse
import os
import queue
import sys
import time
from datetime import datetime

from textual.app import App, ComposeResult
from textual.binding import Binding
from textual.containers import Horizontal
from textual.widgets import Footer, Input, RichLog, Static

from ..assistant.control import default_control_path
from ..config import load_config
from .client import ControlClient
from .dialogs import ConfirmClear
from .state import ViewState


class MainPane(Horizontal):
    """The triage/chat row; allowed to fill the screen when maximized."""
    ALLOW_MAXIMIZE = True


class AssistantTUI(App):
    """Live view of transcriptions, commands, and chat."""

    ALLOW_IN_MAXIMIZED_VIEW = ".__never__"

    CSS = """
    #status { height: 1; }
    #transcript { height: 25%; border: solid $accent; }
    #command { height: 3; border: solid $accent; }
    #chat { height: 1fr; border: solid $accent; }
    #edit { display: none; }
    #main { height: 1fr; }
    #main #triage { width: 1fr; height: 100%; border: solid $accent; }
    #main #chat { width: 2fr; height: 100%; }
    #transcript.-maximized { height: 1fr; }
    """

    BINDINGS = [
        Binding("a", "arm", "arm command"),
        Binding("ctrl+r", "execute", "execute"),
        Binding("escape", "cancel", "cancel"),
        Binding("e", "edit", "edit command"),
        Binding("t", "chat", "chat toggle"),
        Binding("question_mark", "toggle_help_panel", "help"),
        Binding("q", "quit", "quit"),
        Binding("m", "toggle_maximize", "maximize pane", show=False),
        Binding("c", "clear", "clear", show=False),
        Binding("ctrl+s", "screenshot", "save screenshot", show=False),
        Binding("ctrl+p", "command_palette", "instructions", priority=True, show=False),
    ]

    def __init__(self, path: str, triage_enabled: bool = False):
        super().__init__()
        self.path = path
        self.triage_enabled = triage_enabled
        self.state = ViewState()
        self.client = ControlClient(path)
        self.connect_error: str | None = None
        self._events: queue.Queue = queue.Queue()
        self._seen_transcripts = 0
        self._seen_lower = 0
        self._seen_epoch = 0
        self._seen_triage_epoch = 0
        self._timer = None
        self._last_seg_title = ""
        self._last_inf_title = ""

    def compose(self) -> ComposeResult:
        yield Static(id="status")
        yield RichLog(id="transcript", wrap=True, markup=False)
        yield Static(id="command")
        if self.triage_enabled:
            with MainPane(id="main"):
                yield RichLog(id="triage", wrap=True, markup=False)
                yield RichLog(id="chat", wrap=True, markup=False)
        else:
            yield RichLog(id="chat", wrap=True, markup=False)
        yield Input(placeholder="edit the command, then press Enter", id="edit")
        yield Footer()

    def _set_titles(self) -> None:
        self.query_one("#command", Static).border_title = "command buffer"
        if self.triage_enabled:
            self.query_one("#triage", RichLog).border_title = "triage"
        self._update_titles(self.state)

    def _update_titles(self, state) -> None:
        seg_title = f"transcripts  [seg {state.queue_depth}]"
        if seg_title != self._last_seg_title:
            self.query_one("#transcript", RichLog).border_title = seg_title
            self._last_seg_title = seg_title

        inf_suffix = " · busy" if state.inference_in_flight else ""
        inf_title = f"conversation / command output  [inf {state.inference_depth}{inf_suffix}]"
        if inf_title != self._last_inf_title:
            self.query_one("#chat", RichLog).border_title = inf_title
            self._last_inf_title = inf_title

    def on_mount(self) -> None:
        self._set_titles()
        try:
            self.client.connect()
            self.client.subscribe(viewer=True)
        except OSError as e:
            self.connect_error = f"cannot connect to {self.path}: {e}"
            self.exit()
            return
        self.client.start(
            lambda event: self._events.put(event), on_state=self._on_conn_state
        )
        self._timer = self.set_interval(0.1, self._drain)

    def _on_conn_state(self, connected: bool) -> None:
        self.state.connected = connected

    def _drain(self) -> None:
        while True:
            try:
                event = self._events.get_nowait()
            except queue.Empty:
                break
            self.state.ingest(event)
        if self.screen.maximized is None:
            self._restore_panes()
        self._render()

    def _render(self) -> None:
        if not self.query("#status"):
            # A timer tick can land while the app is shutting down, after the
            # widgets are unmounted; there is nothing left to draw into.
            return
        state = self.state
        health = "ok" if state.health_ok() else "DOWN"
        conn = "ok" if state.connected else "LOST"
        speaking = "●" if state.speech_open else "○"
        status = (
            f"[{state.mode_style()} bold]{state.mode}[/]  "
            f"stt:{health}  conn:{conn}  speech:{speaking}  "
            f"worker_stall:{state.worker_stall_s:.1f}s  "
            f"phase:{state.phase}"
        )
        self.query_one("#status", Static).update(status)
        self._update_titles(state)

        edited = " (edited)" if state.command_edited else ""
        self.query_one("#command", Static).update(
            f"command{edited}: {state.command_text or '<empty>'}"
        )

        transcript_log = self.query_one("#transcript", RichLog)
        lines = state.transcript_lines()
        while self._seen_transcripts < len(lines):
            transcript_log.write(lines[self._seen_transcripts])
            self._seen_transcripts += 1

        lower_log = self.query_one("#chat", RichLog)
        if self._seen_epoch != state.chat_epoch:
            lower_log.clear()
            self._seen_lower = 0
            self._seen_epoch = state.chat_epoch
        lower_lines = state.lower_lines()
        while self._seen_lower < len(lower_lines):
            lower_log.write(lower_lines[self._seen_lower])
            self._seen_lower += 1

        if self.triage_enabled:
            triage_log = self.query_one("#triage", RichLog)
            if self._seen_triage_epoch != state.triage_epoch:
                triage_log.clear()
                self._seen_triage_epoch = state.triage_epoch
                triage_log.write("\n".join(state.triage_lines()))

    def _send(self, type_: str, **fields) -> None:
        try:
            self.client.send(type_, **fields)
        except (OSError, RuntimeError) as e:
            self.query_one("#status", Static).update(f"send failed: {e}")

    def action_arm(self) -> None:
        self._send("command_start")

    def action_execute(self) -> None:
        self._send("command_execute")

    def action_cancel(self) -> None:
        self._send("command_cancel")

    def action_chat(self) -> None:
        self._send("chat_toggle")

    def action_clear(self) -> None:
        self.push_screen(ConfirmClear(), self._on_clear_confirmed)

    def _on_clear_confirmed(self, confirmed: bool) -> None:
        if confirmed:
            self._send("clear")

    def action_edit(self) -> None:
        edit = self.query_one("#edit", Input)
        edit.value = self.state.command_text
        edit.display = True
        edit.focus()

    def _active_pane(self):
        focused = self.screen.focused
        if focused is not None:
            for node in focused.ancestors_with_self:
                if getattr(node, "id", None) in ("transcript", "triage", "chat"):
                    if self.query(f"#{node.id}"):
                        return self.query_one(f"#{node.id}")
                    return node
        fallback_id = "triage" if self.triage_enabled else "chat"
        if self.query(f"#{fallback_id}"):
            return self.query_one(f"#{fallback_id}")
        return None

    def _restore_panes(self) -> None:
        if self.query("#triage"):
            self.query_one("#triage").display = True
        if self.query("#chat"):
            self.query_one("#chat").display = True

    def action_toggle_maximize(self) -> None:
        screen = self.screen
        if screen.maximized is not None:
            screen.minimize()
            self._restore_panes()
            return
        pane = self._active_pane()
        if pane is None:
            return
        if pane.id in ("triage", "chat") and self.query("#main"):
            other_id = "#chat" if pane.id == "triage" else "#triage"
            other = self.query_one(other_id) if self.triage_enabled else None
            if other is not None:
                other.display = False
            screen.maximize(self.query_one("#main"))
        else:
            screen.maximize(pane)

    def action_toggle_help_panel(self) -> None:
        from textual.css.query import NoMatches
        from textual.widgets import HelpPanel
        try:
            self.screen.query_one(HelpPanel)
        except NoMatches:
            self.action_show_help_panel()
        else:
            self.action_hide_help_panel()

    def action_screenshot(self) -> None:
        name = "stt-tui-%s.svg" % datetime.now().strftime("%Y%m%d-%H%M%S")
        try:
            save = getattr(self, "save_screenshot", None)
            if callable(save):
                path = save(name)
            else:
                svg = self.export_screenshot()
                path = os.path.join(os.getcwd(), name)
                with open(path, "w", encoding="utf-8") as f:
                    f.write(svg)
        except Exception as e:
            try:
                self.notify("screenshot failed: %s" % e)
            except Exception:
                pass
            return
        try:
            self.notify("saved " + str(path))
        except Exception:
            try:
                self.query_one("#status", Static).update("saved " + str(path))
            except Exception:
                pass

    def on_input_submitted(self, event: Input.Submitted) -> None:
        self._send("command_edit", text=event.value)
        edit = self.query_one("#edit", Input)
        edit.display = False

    def on_unmount(self) -> None:
        if self._timer is not None:
            self._timer.stop()
        self.client.close()


def _dump(path: str, seconds: float) -> int:
    """Headless: print the event stream for a while (verification / debugging)."""
    client = ControlClient(path)
    try:
        client.connect()
    except OSError as e:
        print(f"error: cannot connect to {path}: {e}", file=sys.stderr)
        return 1

    state = ViewState()
    client.subscribe(viewer=True)

    def on_event(event):
        state.ingest(event)
        print(f"{event.get('type'):<12} {event}", flush=True)

    client.start(on_event)
    try:
        time.sleep(seconds)
    except KeyboardInterrupt:
        pass
    finally:
        client.close()
    print(f"mode={state.mode} stt_ok={state.stt_ok} transcripts={len(state.transcripts)}")
    return 0


def main() -> int:
    """Entry point for `stt-tui`."""
    parser = argparse.ArgumentParser(description="Assistant TUI (watch + drive)")
    parser.add_argument("--config", default=None, help="Path to config.toml")
    parser.add_argument("--socket", default=None, help="Override the control socket path")
    parser.add_argument(
        "--dump",
        type=float,
        default=None,
        metavar="SECONDS",
        help="Headless: print the event stream for N seconds and exit",
    )
    args = parser.parse_args()

    cfg = load_config(args.config)
    path = args.socket
    if path is None:
        path = cfg.assistant.control_socket or default_control_path()

    if args.dump is not None:
        return _dump(path, args.dump)

    app = AssistantTUI(path, triage_enabled=cfg.triage.enabled)
    app.run()
    if app.connect_error:
        print(f"error: {app.connect_error}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())