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

"""Pure reduction of assistant/STT events into TUI view state.

Kept free of Textual so the behaviour is unit-testable: attach, mode
self-correction from the heartbeat, command editing, bounded buffers.
"""

from dataclasses import dataclass, field
from datetime import datetime

from rich.text import Text

MAX_TRANSCRIPTS = 200
MAX_CHAT = 200
MAX_LOWER = 500

MODE_STYLES = {
    "idle": "green",
    "command_armed": "yellow",
    "executing": "cyan",
    "chat": "magenta",
}

# Time, language tag, and body text must use disjoint colors so they are never
# confused: these fixed header styles never collide with the body palette.
TIME_STYLE = "bold cyan"
LANG_TAG_STYLE = "bold yellow"
REASON_STYLE = "italic dim"
SPEAKER_NAMED_STYLE = "bold green"
SPEAKER_UNKNOWN_STYLE = "bold red"

LANG_TEXT_STYLES = {
    "es": "green",
    "en": "bright_blue",
    "fr": "magenta",
    "de": "bright_green",
    "pt": "red",
    "it": "bright_magenta",
}
DEFAULT_LANG_TEXT_STYLE = "white"


def _num(value) -> float | None:
    """Best-effort float coercion; None when missing/non-numeric."""
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


def _fmt_clock(ts) -> str:
    """Render a transcript timestamp as local HH:MM:SS (epoch or ISO-8601)."""
    value = _num(ts)
    if value is not None:
        try:
            return datetime.fromtimestamp(value).strftime("%H:%M:%S")
        except (OverflowError, OSError, ValueError):
            return "--:--:--"
    if isinstance(ts, str):
        try:
            return datetime.fromisoformat(ts).strftime("%H:%M:%S")
        except ValueError:
            pass
    return "--:--:--"


def _fmt_gate(item: dict) -> str:
    """Format the gating confidence defensively; '?' when absent."""
    for key in ("answer_confidence", "confidence"):
        value = _num(item.get(key))
        if value is not None:
            return f"{value:.2f}"
    return "?"


@dataclass
class ViewState:
    """Everything the TUI renders."""

    mode: str = "idle"
    stt_ok: bool = False
    device_ok: bool = True
    connected: bool = True
    speech_open: bool = False
    worker_stall_s: float = 0.0
    queue_depth: int = 0
    inference_depth: int = 0
    inference_in_flight: bool = False
    ready: bool = False
    phase: str = "idle"
    command_text: str = ""
    command_edited: bool = False
    in_flight: bool = False
    transcripts: list[dict] = field(default_factory=list)
    transcripts_total: int = 0
    chat: list[dict] = field(default_factory=list)
    lower: list[tuple[str, bool]] = field(default_factory=list)
    tool: dict | None = None
    tool_result: dict | None = None
    last_error: dict | None = None
    chat_epoch: int = 0
    triage: dict | None = None
    triage_state: str = "disabled"
    triage_reason: str = ""
    triage_epoch: int = 0
    triage_ts: float | None = None
    questions: dict | None = None
    questions_epoch: int = 0
    speakers: list[dict] = field(default_factory=list)
    session_dir: str | None = None

    def ingest(self, event: dict) -> None:
        """Fold one event into the view state."""
        etype = event.get("type")

        if etype == "mode":
            self._set_mode(event.get("mode", self.mode))
        elif etype == "status":
            self._set_mode(event.get("mode", self.mode))
            self.ready = bool(event.get("ready", self.ready))
            self.phase = event.get("phase", self.phase)
            self.command_text = event.get("command_text", self.command_text)
            self.stt_ok = bool(event.get("stt_ok", self.stt_ok))
            if "speech_open" in event:
                self.speech_open = bool(event["speech_open"])
            changed = False
            if "triage" in event:
                self.triage = event["triage"]
                changed = True
            if "triage_state" in event:
                self.triage_state = event["triage_state"]
                changed = True
            if "triage_reason" in event:
                self.triage_reason = event["triage_reason"]
                changed = True
            if changed:
                self.triage_epoch += 1
            if "questions" in event:
                # Static config from the analyzer: independent of the triage epoch.
                self.questions = event["questions"]
                self.questions_epoch += 1
            self.session_dir = event.get("session_dir") or self.session_dir
        elif etype == "session":
            self.session_dir = event.get("session_dir") or self.session_dir
        elif etype == "triage":
            stage = event.get("stage")
            if stage == "cleared":
                self.triage = None
                self.triage_ts = None
            else:
                # A triage payload can only come from a ready lane: authoritative.
                self.triage_state = "ok"
                self.triage_reason = ""
                self.triage = event.get("triage")
                self.triage_ts = event.get("ts", self.triage_ts)
            self.triage_epoch += 1
        elif etype == "speech_start":
            self.speech_open = True
        elif etype == "speech_stop":
            self.speech_open = False
        elif etype == "heartbeat":
            # The heartbeat re-states the mode: a missed `mode` event self-corrects.
            self._set_mode(event.get("mode", self.mode))
            self.stt_ok = bool(event.get("stt_ok", self.stt_ok))
            self.in_flight = bool(event.get("in_flight", self.in_flight))
            self.worker_stall_s = float(event.get("worker_stall_s", self.worker_stall_s) or 0.0)
            if "queue_depth" in event:
                self.queue_depth = int(event["queue_depth"] or 0)
            if "inference_depth" in event:
                self.inference_depth = int(event["inference_depth"] or 0)
            if "inference_in_flight" in event:
                self.inference_in_flight = bool(event["inference_in_flight"])
            if "device_ok" in event:
                self.device_ok = bool(event["device_ok"])
        elif etype == "command_text":
            self.command_text = event.get("text", "")
            self.command_edited = bool(event.get("edited", False))
            self.phase = event.get("phase", self.phase)
        elif etype == "speakers":
            self.speakers = list(event.get("speakers") or [])
        elif etype == "transcript":
            self.session_dir = event.get("session_dir") or self.session_dir
            self.transcripts.append(event)
            del self.transcripts[:-MAX_TRANSCRIPTS]
            self.transcripts_total += 1
        elif etype == "chat_message":
            self.chat.append(event)
            del self.chat[:-MAX_CHAT]
            self.append_lower(
                f"{event.get('role', '?')}: {event.get('text', '')}", from_chat=True
            )
        elif etype == "tool_call":
            self.tool = event
            args = event.get("args") or {}
            suffix = f" {args}" if args else ""
            self.append_lower(f"-> tool {event.get('tool')}{suffix}")
        elif etype == "tool_result":
            self.tool_result = event
            verdict = "ok" if event.get("ok") else "failed"
            self.append_lower(f"<- {verdict}: {event.get('result', '')}")
        elif etype == "error":
            self.last_error = event
            self.append_lower(f"! {event.get('code')}: {event.get('message')}")
        elif etype == "cleared":
            self.clear_output()

    def append_lower(self, line: str, from_chat: bool = False) -> None:
        """Append one line to the lower pane log (chat turns, tool calls, results, errors)."""
        self.lower.append((line, from_chat))
        del self.lower[:-MAX_LOWER]

    def clear_chat(self) -> None:
        """Forget chat turns and their rendered lines, keeping tool/error output."""
        self.chat = []
        self.lower = [(line, from_chat) for line, from_chat in self.lower if not from_chat]
        self.chat_epoch += 1

    def clear_output(self) -> None:
        """Clear conversation/command output: chat, tool calls/results, errors."""
        self.chat        = []
        self.lower       = []
        self.tool        = None
        self.tool_result = None
        self.last_error  = None
        self.triage      = None
        self.triage_ts   = None
        self.chat_epoch += 1
        self.triage_epoch += 1

    def _set_mode(self, mode: str) -> None:
        if self.mode == "chat" and mode != "chat":
            self.clear_chat()
        self.mode = mode

    def lower_lines(self) -> list[str]:
        return [line for line, _ in self.lower]

    def health_ok(self) -> bool:
        """False when the header should go red."""
        return self.stt_ok and self.device_ok

    def mode_style(self) -> str:
        return MODE_STYLES.get(self.mode, "white")

    def transcript_lines(self) -> list[Text]:
        """Render each transcript as a styled two-line block: header then text."""
        blocks: list[Text] = []
        for e in self.transcripts:
            lang = str(e.get("lang") or "?").lower()
            text_style = LANG_TEXT_STYLES.get(lang, DEFAULT_LANG_TEXT_STYLE)
            reason = str(e.get("end_reason") or "").strip()
            duration = _num(e.get("duration_s"))
            speaker = str(e.get("speaker") or "Unknown")
            speaker_style = (
                SPEAKER_NAMED_STYLE
                if e.get("speaker_status") == "matched" and speaker not in ("", "Unknown")
                else SPEAKER_UNKNOWN_STYLE
            )

            block = Text()
            block.append(_fmt_clock(e.get("ts")), style=TIME_STYLE)
            block.append("  ")
            block.append(lang.upper(), style=LANG_TAG_STYLE)
            block.append("  ")
            block.append(speaker, style=speaker_style)
            block.append("  ")
            block.append(f"{duration:.1f}s" if duration is not None else "?", style="dim")
            if reason:
                block.append("  ")
                block.append(reason, style=REASON_STYLE)
            block.append("\n")
            block.append("  " + str(e.get("text", "")), style=text_style)
            blocks.append(block)
        return blocks

    def chat_lines(self) -> list[str]:
        return [f"{m.get('role', '?')}: {m.get('text', '')}" for m in self.chat]

    def triage_renderables(self) -> list[Text]:
        """triage_lines() as Text, with the last-triage clock header when available."""
        out = [Text(line) for line in self.triage_lines()]
        if (
            self.triage_ts is not None
            and self.triage_state == "ok"
            and isinstance(self.triage, dict)
        ):
            clock = Text()
            clock.append(_fmt_clock(self.triage_ts), style=TIME_STYLE)
            out.insert(0, clock)
        return out

    def triage_lines(self) -> list[str]:
        """Render the triage panel lines. Pure and defensive: never raises."""
        if self.triage_state != "ok":
            reason = self.triage_reason or self.triage_state or "unavailable"
            return [f"triage: unavailable — {reason}"]
        if not isinstance(self.triage, dict):
            return ["triage: unavailable — no data"]

        order = self.triage.get("order")
        items = self.triage.get("items")
        lines: list[str] = []
        if isinstance(order, list) and isinstance(items, dict):
            for qid in order:
                item = items.get(qid)
                if not isinstance(item, dict):
                    lines.append(f"{qid} ⚠" if item else str(qid))
                    continue
                kind = item.get("kind")
                flagged = bool(item.get("flagged"))
                mark = " ⚠" if flagged else ""
                if kind == "choice":
                    line = f"{qid}: {item.get('choice', '?')} ({_fmt_gate(item)}){mark}"
                elif kind == "noul":
                    p = _num(item.get("p"))
                    bar = "█" * int(p * 10) if p is not None else ""
                    bar = bar.ljust(10)
                    ptxt = f"{p:.2f}" if p is not None else "?"
                    line = f"{bar} {qid}: {ptxt}{mark}"
                elif kind == "score":
                    level = _num(item.get("level"))
                    bar = ("█" * int(level)).ljust(10) if level is not None else " " * 10
                    label = item.get("label", "?")
                    line = f"{bar} {qid}: {label}{mark}"
                else:
                    line = f"{qid}{mark}" if flagged else str(qid)
                lines.append(line)

        brief = self.triage.get("brief")
        if isinstance(brief, str) and brief:
            lines.append("")
            lines.extend(brief.splitlines())

        if not lines:
            return ["triage: no signals"]
        return lines

    def questions_lines(self) -> list[str]:
        """Render the loaded questions as structured lines. Pure and defensive: never raises."""
        entries = self._question_entries()
        if not entries:
            reason = self.triage_reason or self.triage_state or "no questions loaded"
            return [f"questions: unavailable — {reason}"]
        lines: list[str] = []
        for qid, qtype, body in entries:
            lines.append(f"{qid}  [{qtype}]")
            lines.extend(body)
        return lines

    def questions_renderables(self) -> list[Text]:
        """questions_lines() as Text, with the qid header styled (bold qid, dim type)."""
        entries = self._question_entries()
        if not entries:
            return [Text(line) for line in self.questions_lines()]
        out: list[Text] = []
        for qid, qtype, body in entries:
            header = Text()
            header.append(qid, style="bold")
            header.append("  ")
            header.append(f"[{qtype}]", style="dim")
            out.append(header)
            for line in body:
                out.append(Text(line))
        return out

    def _question_entries(self) -> list[tuple[str, str, list[str]]]:
        """(qid, type, body lines) per usable question, in file order. Never raises."""
        entries: list[tuple[str, str, list[str]]] = []
        if not (isinstance(self.questions, dict) and self.questions):
            return entries
        for qid, qdef in self.questions.items():
            if not isinstance(qdef, dict):
                continue
            qtype = str(qdef.get("type") or "?")
            body: list[str] = []
            instructions = qdef.get("instructions")
            if isinstance(instructions, str) and instructions:
                body.append(f"  {instructions}")
            body.extend(self._criteria_lines(qtype, qdef.get("criteria")))
            entries.append((str(qid), qtype, body))
        return entries

    @staticmethod
    def _criteria_lines(qtype: str, criteria) -> list[str]:
        """One line per criterion for a known type; malformed criteria degrade to none."""
        lines: list[str] = []
        if qtype == "choice":
            if isinstance(criteria, dict):
                for label, description in criteria.items():
                    lines.append(f"    {label}  — {description}")
            elif isinstance(criteria, list):
                for label in criteria:
                    lines.append(f"    {label}")
        elif qtype == "score":
            if isinstance(criteria, list):
                for index, level in enumerate(criteria, start=1):
                    lines.append(f"    {index}. {level}")
        elif qtype == "noul":
            if isinstance(criteria, dict):
                for value, description in criteria.items():
                    lines.append(f"    {value}  — {description}")
        return lines