"""Custom Textual widgets for SneakyCode TUI.""" from __future__ import annotations from textual.app import ComposeResult from textual.containers import Horizontal, Vertical from textual.events import Key from textual.screen import ModalScreen from textual.timer import Timer from textual.widgets import Button, Input, Static from rich.text import Text # --------------------------------------------------------------------------- # Modal Dialogs # --------------------------------------------------------------------------- class PermissionModal(ModalScreen[bool]): """Modal dialog for tool permission approval.""" BINDINGS = [("y", "allow", "Allow"), ("n", "deny", "Deny")] def __init__(self, tool_name: str, description: str) -> None: super().__init__() self._tool_name = tool_name self._description = description def compose(self) -> ComposeResult: with Vertical(id="permission-dialog"): yield Static(f"Tool: {self._tool_name}", classes="modal-title") yield Static(self._description, classes="modal-body") with Horizontal(classes="modal-buttons"): yield Button("Allow (y)", id="allow", variant="success") yield Button("Deny (n)", id="deny", variant="error") def on_button_pressed(self, event: Button.Pressed) -> None: self.dismiss(event.button.id == "allow") def action_allow(self) -> None: self.dismiss(True) def action_deny(self) -> None: self.dismiss(False) class SessionResumeModal(ModalScreen[bool]): """Modal dialog for session resume prompt.""" BINDINGS = [("y", "resume", "Resume"), ("n", "fresh", "Start Fresh")] def __init__(self, msg_count: int) -> None: super().__init__() self._msg_count = msg_count def compose(self) -> ComposeResult: with Vertical(id="permission-dialog"): yield Static("Resume Session?", classes="modal-title") yield Static( f"Found previous session with {self._msg_count} messages.", classes="modal-body", ) with Horizontal(classes="modal-buttons"): yield Button("Resume (y)", id="resume", variant="success") yield Button("Start Fresh (n)", id="fresh", variant="warning") def on_button_pressed(self, event: Button.Pressed) -> None: self.dismiss(event.button.id == "resume") def action_resume(self) -> None: self.dismiss(True) def action_fresh(self) -> None: self.dismiss(False) # --------------------------------------------------------------------------- # Input with History # --------------------------------------------------------------------------- class HistoryInput(Input): """Input with up/down arrow history cycling.""" def __init__(self, **kwargs: object) -> None: super().__init__(**kwargs) self._history: list[str] = [] self._history_index: int = -1 self._draft: str = "" def record(self, value: str) -> None: """Record a submitted value in history.""" if value.strip() and (not self._history or self._history[-1] != value): self._history.append(value) self._history_index = -1 self._draft = "" def on_key(self, event: Key) -> None: if event.key == "up" and self._history: event.prevent_default() if self._history_index == -1: self._draft = self.value self._history_index = len(self._history) - 1 elif self._history_index > 0: self._history_index -= 1 self.value = self._history[self._history_index] self.cursor_position = len(self.value) elif event.key == "down" and self._history_index != -1: event.prevent_default() self._history_index += 1 if self._history_index >= len(self._history): self._history_index = -1 self.value = self._draft else: self.value = self._history[self._history_index] self.cursor_position = len(self.value) # --------------------------------------------------------------------------- # Status Bar # --------------------------------------------------------------------------- class StatusBar(Static): """Single-line status bar showing token usage, iteration count, and streaming spinner.""" _SPINNER = "\u280b\u2819\u2839\u2838\u283c\u2834\u2826\u2827\u2807\u280f" DEFAULT_CSS = """ StatusBar { dock: bottom; height: 1; background: $surface; color: $text-muted; padding: 0 2; } """ def __init__(self) -> None: super().__init__("") self._tokens: int = 0 self._budget: int = 0 self._iteration: int = 0 self._max_iterations: int = 0 self._streaming: bool = False self._spinner_frame: int = 0 self._spinner_timer: Timer | None = None self._stream_tokens: int = 0 self._active_skill: str | None = None def update_tokens(self, tokens: int, budget: int) -> None: """Update the token usage display.""" self._tokens = tokens self._budget = budget self._refresh_display() def update_iteration(self, iteration: int, max_iterations: int) -> None: """Update the iteration count display.""" self._iteration = iteration self._max_iterations = max_iterations self._refresh_display() def start_streaming(self) -> None: """Start the animated thinking spinner.""" self._streaming = True self._spinner_frame = 0 self._stream_tokens = 0 self._spinner_timer = self.set_interval(0.08, self._tick_spinner) self._refresh_display() def stop_streaming(self) -> None: """Stop the animated thinking spinner.""" self._streaming = False if self._spinner_timer: self._spinner_timer.stop() self._spinner_timer = None self._refresh_display() def update_stream_tokens(self, tokens: int) -> None: """Update estimated token count during streaming.""" self._stream_tokens = tokens def _tick_spinner(self) -> None: self._spinner_frame = (self._spinner_frame + 1) % len(self._SPINNER) self._refresh_display() def set_active_skill(self, skill_name: str | None) -> None: """Set or clear the active skill indicator.""" self._active_skill = skill_name self._refresh_display() def _refresh_display(self) -> None: """Rebuild the status bar text.""" parts: list[str] = [] if self._active_skill: parts.append(f"[Skill: {self._active_skill}]") if self._streaming: spinner = self._SPINNER[self._spinner_frame] parts.append(f"{spinner} Thinking") if self._stream_tokens > 0: parts.append(f"~{self._stream_tokens:,} tokens") if self._budget > 0: parts.append(f"Tokens: ~{self._tokens:,} / {self._budget:,}") if self._max_iterations > 0: parts.append(f"Iteration {self._iteration}/{self._max_iterations}") self.update(Text(" \u2502 ".join(parts), style="dim")) # --------------------------------------------------------------------------- # Streaming Display # --------------------------------------------------------------------------- class StreamingStatic(Static): """A Static widget that stays mounted but hidden during non-streaming periods. During streaming, call show() to make visible and update() with partial content. When streaming ends, call hide() to conceal. """ def show_streaming(self) -> None: """Make the widget visible for streaming.""" self.add_class("visible") def hide_streaming(self) -> None: """Hide the widget and clear content.""" self.remove_class("visible") self.update("")