display.py

"""Terminal presentation: ANSI colors, the startup banner, and the animated
status spinner. Nothing in here touches Session state or agent internals —
it only renders strings it's handed."""

import os
import re
import select
import shutil
import sys
import threading
import time

try:
    import termios
    import tty
    _RAW_TERMINAL_SUPPORTED = True
except ImportError:
    # Windows has neither module; arrow-key selection falls back to plain
    # numbered/typed input (see select_one() below).
    _RAW_TERMINAL_SUPPORTED = False

C_AI = "\033[38;2;0;255;180m"   # Mint
C_USER = "\033[96m"             # Cyan
C_DIM = "\033[2m"               # Faded gray
C_SUCCESS = "\033[92m"          # Green
C_WARN = "\033[93m"             # Yellow
C_BLUE = "\033[94m"             # Blue
C_BOLD = "\033[1m"
C_RESET = "\033[0m"
CLEAR_LINE = "\r\033[K"
_ANSI_RE = re.compile(r"\033\[[0-9;]*m")

# How long to wait for the rest of an escape sequence before treating the
# Escape byte as the user pressing Esc. Terminals send the whole sequence in
# one burst, so anything still unread after this is not part of one.
_ESCAPE_TIMEOUT = 0.05

_BANNER = f"""{C_AI}{C_BOLD}
 ████████╗██╗   ██╗███████╗███████╗██╗   ██╗
 ╚══██╔══╝██║   ██║██╔════╝██╔════╝╚██╗ ██╔╝
    ██║   ██║   ██║█████╗  █████╗   ╚████╔╝
    ██║   ██║   ██║██╔══╝  ██╔══╝    ╚██╔╝
    ██║   ╚██████╔╝██║     ██║        ██║
    ╚═╝    ╚═════╝ ╚═╝     ╚═╝        ╚═╝
{C_RESET}"""


def print_logo():
    """Just the ASCII art — no model/status text baked into it, so startup
    info (which model loaded, whether vision is on) always reads as
    something that happened *after* the banner, not part of it."""
    print(_BANNER)


# --- block layout ------------------------------------------------------
#
# Every slash command prints one or more "blocks": an optional heading, then
# either name/description rows (a list of things you can type) or label/value
# rows (details of one thing). The helpers below are the only place that
# decides padding, separators, and color, so /help, /models, /tools, /skills,
# /mcp, /status, /memory and /network info all come out looking like one
# program instead of nine.
#
# Color roles, applied consistently by these helpers:
#   blue + bold   section headings
#   green         identifiers the user can type (commands, model ids, tool
#                 names) and success confirmations
#   yellow        warnings, usage errors, unavailable states
#   dim           labels, hints, and secondary text
#   default       descriptions and values
#
# Names are padded to a width the caller passes in, so a caller with several
# blocks can align them all to one shared column (see column_width) rather
# than letting each block pick its own and look ragged next to the others.

_GUTTER = "  "   # separates a padded name/label from its description/value


def column_width(names) -> int:
    """The padded name-column width for a set of names, computed on visible
    length (ANSI codes are zero-width but count in len())."""
    return max((len(_ANSI_RE.sub("", n)) for n in names), default=0)


def print_heading(title: str, note: str = ""):
    """A block heading, always preceded by one blank line so consecutive
    blocks are separated identically no matter which command printed them."""
    suffix = f"{_GUTTER}{C_DIM}{note}{C_RESET}" if note else ""
    print(f"\n{C_BOLD}{C_BLUE}{title}{C_RESET}{suffix}")


def print_rows(rows, width: int = None, name_color: str = C_SUCCESS, indent: str = "  "):
    """Name/description rows: `  <name padded>  <description>`.

    `rows` is an iterable of (name, description). No bullet, no dash - the
    padded column is what separates the two halves, and a dash only adds a
    second, redundant separator that shifts with every name length.
    """
    rows = list(rows)
    if width is None:
        width = column_width(name for name, _ in rows)
    for name, description in rows:
        padded = name + " " * (width - len(_ANSI_RE.sub("", name)))
        print(f"{indent}{name_color}{padded}{C_RESET}{_GUTTER}{description}")


def print_fields(fields, width: int = None, indent: str = "  "):
    """Label/value rows: `  <label padded> : <value>`, for the details of one
    thing (a model card, session status). `fields` is an iterable of
    (label, value); the colon sits in one column for the whole block."""
    fields = list(fields)
    if width is None:
        width = column_width(label for label, _ in fields)
    for label, value in fields:
        print(f"{indent}{C_DIM}{label.ljust(width)}{C_RESET} : {value}")


def rate_limits_line(limits: dict) -> str:
    """Renders whichever published rate limits a model card actually carries.

    Providers don't all publish the same set — some quote per-minute limits
    only, some add daily caps — so this formats the keys that are present
    instead of requiring all four and raising KeyError on a card with fewer.
    Lives here, next to the other formatting, because both the startup block
    and '/models info' render the same figures and had drifted into two
    copies of this loop."""
    labels = [
        ("requests_per_minute", "req/min"),
        ("requests_per_day", "req/day"),
        ("tokens_per_minute", "tok/min"),
        ("tokens_per_day", "tok/day"),
    ]
    parts = [f"{limits[key]:,} {unit}" for key, unit in labels if limits.get(key)]
    return ", ".join(parts)


def print_session_info(model_card: dict, vision: bool, network_mode: str):
    """Printed once, after the logo and after the model has finished loading:
    what this session is running, as one key/value block in the same shape
    every command uses, then a pointer to /help.

    One block rather than the previous run of prose lines ("Starting in X
    mode.", "Ready: <model> (vision on)", a bare limits line, "Active model:
    <model> (vision)") — those repeated the model name twice, put the network
    mode a blank line away from everything it affects, and left the reader
    scanning four different sentence shapes for four facts."""
    fields = [
        ("mode", network_mode),
        ("model", model_card["name"]),
        ("provider", model_card["provider"]),
        ("vision", "on" if vision else "off"),
    ]
    if model_card.get("context_length"):
        fields.append(("context", f"{model_card['context_length']:,} tok"))
    limits = rate_limits_line(model_card.get("rate_limits") or {})
    if limits:
        fields.append(("rate limits", limits))
    print_fields(fields)
    print(f"{C_DIM}Type /help to see everything Tuffy can do.{C_RESET}\n")


class Spinner:
    """Terminal spinner for AI status updates."""

    MAX_LABEL = 64

    def __init__(self, label: str = "thinking"):
        self.label = label
        self._stop_event = threading.Event()
        self._thread = None
        self._lock = threading.Lock()
        self._last_rows = 0  # terminal rows the last-drawn frame wrapped onto

    def set_label(self, label: str):
        label = " ".join(str(label).split())

        if len(label) > self.MAX_LABEL:
            label = label[: self.MAX_LABEL - 1] + "…"

        with self._lock:
            self.label = label or "thinking"

    def _clear_last_render(self):
        """Erases every terminal row the previous frame drew on, not just
        the current one. A long label can push 'AI ❯ label...' past the
        terminal width and wrap onto a second row; \\r\\033[K only clears
        the row the cursor is on, so a naive clear leaves the wrapped-over
        remainder (including stray '...' dots) stuck in the scrollback."""
        if self._last_rows > 1:
            sys.stdout.write(f"\033[{self._last_rows - 1}A")
        sys.stdout.write(CLEAR_LINE)
        for _ in range(self._last_rows - 1):
            sys.stdout.write("\033[B\033[K")
        if self._last_rows > 1:
            sys.stdout.write(f"\033[{self._last_rows - 1}A")
        self._last_rows = 0

    def _render(self, text: str):
        self._clear_last_render()
        sys.stdout.write(text)
        sys.stdout.flush()

        cols = shutil.get_terminal_size(fallback=(80, 24)).columns
        visible_len = len(_ANSI_RE.sub("", text))
        self._last_rows = max(1, -(-visible_len // cols))  # ceil div

    def start(self):
        if self._thread is not None:
            return

        self._stop_event.clear()

        sys.stdout.write("\033[?25l")
        sys.stdout.flush()

        def run():
            frames = ["", ".", "..", "..."]

            i = 0
            while not self._stop_event.is_set():
                with self._lock:
                    label = self.label

                self._render(
                    f"{C_AI}AI ❯{C_RESET} "
                    f"{C_DIM}{label}{frames[i % len(frames)]}{C_RESET}"
                )

                i += 1
                time.sleep(0.4)

        self._thread = threading.Thread(target=run, daemon=True)
        self._thread.start()

    def stop(self, show_prompt: bool = True):
        if self._thread is None:
            return

        self._stop_event.set()
        self._thread.join()

        self._thread = None

        sys.stdout.write("\033[?25h")
        sys.stdout.flush()

        if show_prompt:
            self._render(f"{C_AI}AI ❯{C_RESET} ")
        else:
            self._clear_last_render()
            sys.stdout.flush()


def _read_key() -> str:
    """Reads one keypress from stdin (already in raw mode - see select_one,
    which sets it once for the whole selection loop rather than per key),
    resolving arrow-key escape sequences ('\\x1b[A'/'\\x1b[B') to 'up'/'down'
    and a bare Escape to 'cancel'. Falls back to the raw character for
    anything else (Enter, Ctrl-C, plain letters).

    A bare Esc and an arrow key start with the same byte, so the two are told
    apart by whether anything follows: with raw mode's VMIN/VTIME left at the
    defaults a lone Esc would block the next read forever waiting for the rest
    of a sequence that never comes, so the continuation bytes are read only if
    select() says they are already there.

    Reads go through os.read on the raw file descriptor, never
    sys.stdin.read: sys.stdin is buffered, so reading one character off it
    pulls the whole available block into Python's own buffer, leaving the fd
    empty. select() then reports "nothing pending" for an arrow key whose
    '[B' tail is already buffered in userspace, and every arrow press would be
    treated as Esc."""
    fd = sys.stdin.fileno()
    ch = os.read(fd, 1).decode(errors="ignore")
    if ch == "\x1b":
        ready, _, _ = select.select([fd], [], [], _ESCAPE_TIMEOUT)
        if not ready:
            return "cancel"
        ch += os.read(fd, 2).decode(errors="ignore")

    if ch == "\x1b[A":
        return "up"
    if ch == "\x1b[B":
        return "down"
    if ch in ("\r", "\n"):
        return "enter"
    if ch in ("\x03", "\x04"):  # Ctrl-C / Ctrl-D
        return "interrupt"
    return ch


def select_one(prompt: str, options: list[str], current: str, labels: list[str] = None) -> str:
    """Interactive arrow-key selector: Up/Down moves the cursor, Enter
    confirms, Esc (or Ctrl-C) backs out. Starts on 'current' so the prompt
    always shows what's active right now. Falls back to plain typed input
    (Enter = keep current) on platforms without a raw terminal (e.g. Windows,
    or stdin not a TTY — piped input in tests/scripts).

    Esc raises KeyboardInterrupt, the same exit Ctrl-C already took, so a
    caller has one "the user backed out" path to handle rather than a
    sentinel return value that every call site would have to remember to
    check against a real choice.

    `options` are the values returned; `labels`, when given, are what is shown
    for each one (same length, same order). That split is what lets a caller
    select on a short value (a model id) while displaying a whole row of
    columns for it - the selector never has to know what the columns mean.
    """
    shown = labels if labels is not None else options
    if not _RAW_TERMINAL_SUPPORTED or not sys.stdin.isatty():
        # Non-interactive: print the choices before asking, since the
        # caller's list may be the only place they are written down. The
        # prompt itself is not repeated above the list - only the input line
        # carries it, with the current value as the Enter-to-keep default.
        for option, label in zip(options, shown):
            marker = f"{C_SUCCESS}*{C_RESET}" if option == current else " "
            print(f"  {marker} {label}")
        raw = input(f"{C_DIM}{prompt}{C_RESET} [{current}] \u276f ").strip().lower()
        for opt in options:
            if raw == opt.lower():
                return opt
        return current

    index = options.index(current) if current in options else 0

    def render():
        # \r\n, not \n: stdin is in raw mode for the duration of the
        # selection loop, which also disables ONLCR output translation, so a
        # bare \n would move the cursor down without returning to column 0.
        sys.stdout.write(f"{C_DIM}{prompt}{C_RESET}\r\n")
        for i, (option, label) in enumerate(zip(options, shown)):
            selected = i == index
            marker = f"{C_SUCCESS}{C_BOLD}❯{C_RESET}" if selected else " "
            # Selected row: green + bold. The contrast comes from every other
            # row being dimmed - bold alone, among full-brightness rows, was
            # barely distinguishable.
            row = (
                f"{C_SUCCESS}{C_BOLD}{label}{C_RESET}" if selected
                else f"{C_DIM}{label}{C_RESET}"
            )
            current_tag = f" {C_DIM}(current){C_RESET}" if option == current else ""
            sys.stdout.write(f"  {marker} {row}{current_tag}\r\n")
        sys.stdout.flush()

    def clear(n_lines: int):
        sys.stdout.write(f"\033[{n_lines}A")
        for _ in range(n_lines):
            sys.stdout.write("\033[K\033[B")
        sys.stdout.write(f"\033[{n_lines}A")
        sys.stdout.flush()

    n_lines = len(options) + 1
    fd = sys.stdin.fileno()
    old_settings = termios.tcgetattr(fd)
    sys.stdout.write("\033[?25l")  # hide cursor
    render()
    try:
        tty.setraw(fd)
        while True:
            key = _read_key()
            if key == "up":
                index = (index - 1) % len(options)
            elif key == "down":
                index = (index + 1) % len(options)
            elif key == "enter":
                break
            elif key in ("interrupt", "cancel"):
                clear(n_lines)
                sys.stdout.write("\033[?25h")
                sys.stdout.flush()
                raise KeyboardInterrupt
            else:
                continue
            clear(n_lines)
            render()
    finally:
        termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
        sys.stdout.write("\033[?25h")  # show cursor

    # No "prompt: answer" echo on the way out: every caller already announces
    # what it did with the value ("Switched to model X", "Starting in offline
    # mode"), so echoing the choice here only duplicated the next line.
    clear(n_lines)
    sys.stdout.flush()
    return options[index]