path: root/tools/tui.py

#!/usr/bin/env python3
"""Minimal curses TUI widgets for the interactive tools.

One screen per decision: a directory browser, an expandable checkbox
tree, a single-line text editor, a single-choice menu, and a yes/no
confirm. There is no framework — every widget is a function that runs
its own key loop on a curses window and returns the chosen value.

Common key bindings:

  Up/Down (or k/j)   move the cursor
  Enter              accept
  Esc                abort the whole wizard (raises WizardCancelled)

On screens without typed text (menus, confirm, tree, browser) 'q' also
aborts; inside text editors it is an ordinary character.
"""

import os
from pathlib import Path
from typing import Callable, List, Optional, Sequence, Tuple

# Make Esc register quickly instead of pausing for an escape sequence.
os.environ.setdefault("ESCDELAY", "25")


class WizardCancelled(Exception):
    """Raised when the user presses Esc to abort the wizard."""


# ---------------------------------------------------------------------------
# Shared drawing helpers
# ---------------------------------------------------------------------------

def _addstr(scr, y: int, x: int, text: str, attr: int = 0) -> None:
    """addstr that ignores out-of-bounds and terminal-capability errors."""
    try:
        scr.addstr(y, x, text, attr)
    except Exception:
        pass


def _fit(text: str, width: int) -> str:
    """Truncate TEXT to WIDTH columns, appending '~' when cut."""
    if width < 1:
        return ""
    if len(text) <= width:
        return text
    return text[: max(0, width - 1)] + "~"


class Frame:
    """A screen frame: title, scrolling body rows, message and footer.

    Widgets append styled body rows via mark(), call draw() after every
    state change, and read keys through get_key()/edit_line().
    """

    def __init__(self, scr, title: str, footer: str):
        import curses
        self.curses = curses
        self.scr = scr
        self.title = title
        self.footer = footer
        self.message = ""            # transient status line
        self.message_attr = None     # None -> bold reverse video
        self.rows: List[dict] = []   # {text, attr, indent}
        self.scroll = 0
        self.cursor = 0              # highlighted row index

    def mark(self, text: str, attr: int = 0, indent: int = 0) -> None:
        self.rows.append({"text": text, "attr": attr, "indent": indent})

    def draw(self) -> None:
        curses = self.curses
        scr = self.scr
        scr.erase()
        height, width = scr.getmaxyx()
        if height < 6 or width < 20:
            _addstr(scr, 0, 0, _fit("Terminal too small", width - 1),
                    curses.A_BOLD)
            scr.refresh()
            return
        top = 2
        visible = height - 3 - top
        if visible < 1:
            visible = 1
        # Keep the cursor inside the viewport.
        if self.cursor < self.scroll:
            self.scroll = self.cursor
        elif self.cursor >= self.scroll + visible:
            self.scroll = self.cursor - visible + 1
        if self.scroll + visible > len(self.rows):
            self.scroll = max(0, len(self.rows) - visible)
        scrolling = len(self.rows) > visible
        indicator = f" {self.cursor + 1}/{len(self.rows)} " if scrolling else ""
        title_width = width - 1 - (len(indicator) if indicator else 0)
        _addstr(scr, 0, 0, _fit(self.title, title_width),
                curses.A_BOLD | curses.A_UNDERLINE)
        for index in range(self.scroll,
                           min(len(self.rows), self.scroll + visible)):
            row = self.rows[index]
            line = "  " * row["indent"] + row["text"]
            attr = row["attr"]
            if index == self.cursor:
                attr |= curses.A_REVERSE
            _addstr(scr, top + index - self.scroll, 0,
                    _fit(line, width - 1), attr)
        if indicator:
            _addstr(scr, 0, max(0, width - len(indicator)), indicator,
                    curses.A_DIM)
        if self.message:
            attr = self.message_attr
            if attr is None:
                attr = curses.A_BOLD | curses.A_REVERSE
            _addstr(scr, height - 2, 0, _fit(self.message, width - 1), attr)
        _addstr(scr, height - 1, 0, _fit(self.footer, width - 1), curses.A_DIM)
        scr.refresh()

    # -- key helpers ------------------------------------------------------

    def get_key(self, cancel_keys: Sequence[int] = (27,)) -> int:
        """Read one key; cancel keys and Ctrl-C raise WizardCancelled."""
        try:
            key = self.scr.getch()
        except KeyboardInterrupt:
            raise WizardCancelled() from None
        if key == 3:  # Ctrl-C
            raise WizardCancelled()
        if key in cancel_keys:
            raise WizardCancelled()
        return key

    def edit_line(self, start: str, prompt: str = ""
                  ) -> Optional[str]:
        """Run an inline editor on the message line.

        Returns the edited string on Enter, or None when the user backs
        out with Esc (the caller decides what that means).
        """
        curses = self.curses
        text = start
        while True:
            height, width = self.scr.getmaxyx()
            self.message = ""
            self.draw()
            room = max(1, width - 2 - len(prompt))
            shown = text if len(text) < room else ">" + text[-(room - 2):]
            _addstr(self.scr, height - 2, 0,
                    _fit(f"{prompt}{shown}_", width - 1), curses.A_BOLD)
            self.scr.refresh()
            try:
                key = self.scr.getch()
            except KeyboardInterrupt:
                raise WizardCancelled() from None
            if key == 27:
                return None
            if key in (10, 13):  # Enter
                return text
            if key in (curses.KEY_BACKSPACE, 8, 127):
                text = text[:-1]
            elif 32 <= key < 127:
                text += chr(key)


# ---------------------------------------------------------------------------
# Widget: yes/no confirm
# ---------------------------------------------------------------------------

def confirm(scr, question: str, default: bool = False,
            body: Optional[Sequence[str]] = None) -> bool:
    """Ask a yes/no QUESTION; Enter takes DEFAULT, Esc aborts.

    BODY lines are shown above the question (a summary, for example).
    """
    frame = Frame(scr, question,
                  "y = yes  n = no  Enter = default  Esc = cancel")
    cancel = (27, ord("q"))
    while True:
        frame.rows = []
        for line in body or []:
            frame.mark(line)
        if body:
            frame.mark("")
        hint = "[Y/n]" if default else "[y/N]"
        frame.mark(f"{question} {hint}")
        frame.cursor = len(frame.rows) - 1
        frame.draw()
        key = frame.get_key(cancel)
        if key in (ord("y"), ord("Y")):
            return True
        if key in (ord("n"), ord("N")):
            return False
        if key in (10, 13):
            return default


# ---------------------------------------------------------------------------
# Widget: single-choice menu
# ---------------------------------------------------------------------------

def menu(scr, title: str, options: Sequence[tuple], default_index: int = 0):
    """Show OPTIONS as (label, value) pairs; return the chosen value.

    The cursor starts on DEFAULT_INDEX; Enter returns the highlighted
    option's value.
    """
    frame = Frame(scr, title,
                  "Up/Down = move  Enter = select  Esc = cancel")
    cancel = (27, ord("q"))
    cursor = max(0, min(default_index, len(options) - 1))
    while True:
        frame.rows = []
        for label, _ in options:
            frame.mark(label)
        frame.cursor = cursor
        frame.draw()
        curses = frame.curses
        key = frame.get_key(cancel)
        if key in (curses.KEY_UP, ord("k")):
            cursor = (cursor - 1) % len(options)
        elif key in (curses.KEY_DOWN, ord("j")):
            cursor = (cursor + 1) % len(options)
        elif key in (10, 13):
            return options[cursor][1]


# ---------------------------------------------------------------------------
# Widget: single-line text editor
# ---------------------------------------------------------------------------

def line_edit(scr, title: str, default: str,
              validate: Optional[Callable[[str], Optional[str]]] = None
              ) -> str:
    """Edit one line of text, pre-filled with DEFAULT; Enter accepts.

    VALIDATE receives the entered string and returns an error message or
    None; Enter on an invalid value shows the message and keeps editing.
    Esc aborts the wizard ('q' is an ordinary character here).
    """
    frame = Frame(scr, title,
                  "type to edit  Backspace = erase  Enter = accept  "
                  "Esc = cancel")
    text = default
    error = ""
    while True:
        frame.rows = []
        frame.mark("")
        frame.mark(f"  {text}_")
        frame.cursor = 1
        frame.message = error
        frame.draw()
        curses = frame.curses
        key = frame.get_key()  # Esc only; 'q' must stay typeable
        if key in (10, 13):
            if validate is None:
                return text
            error = validate(text)
            if error is None:
                return text
            error = f"{error}  (edit, then Enter)"
            continue
        if key in (curses.KEY_BACKSPACE, 8, 127):
            text = text[:-1]
        elif 32 <= key < 127:
            text += chr(key)


# ---------------------------------------------------------------------------
# Widget: directory browser
# ---------------------------------------------------------------------------

def _list_dirs(path: Path) -> List[Path]:
    """Return the subdirectories of PATH, sorted, dot-dirs excluded."""
    try:
        entries = [child for child in path.iterdir()
                   if child.is_dir() and not child.name.startswith(".")]
    except OSError:
        return []
    return sorted(entries, key=lambda child: child.name.lower())


def browse_directory(scr, title: str,
                     validate: Optional[Callable[[Path], Optional[str]]] = None,
                     start: Optional[Path] = None
                     ) -> Path:
    """Pick a directory; Enter accepts the directory being listed.

    Right (or l) descends into the highlighted entry, Left/Backspace/u
    goes to the parent, and e edits the path directly. VALIDATE receives
    the listed directory and returns an error message or None; Enter on
    an invalid directory is refused with that message. Esc aborts the
    wizard.
    """
    footer = ("Up/Down = move  Right = open  Left = parent  e = edit "
              "path  Enter = choose this directory  Esc = cancel")
    frame = Frame(scr, title, footer)
    cancel = (27, ord("q"))
    current = Path(start) if start is not None else Path.cwd()
    try:
        current = current.resolve()
    except OSError:
        current = Path.cwd()
    cursor = 0

    def validation_error() -> Optional[str]:
        if validate is None:
            return None
        return validate(current)

    while True:
        entries = _list_dirs(current)
        cursor = max(0, min(cursor, max(0, len(entries) - 1)))
        frame.rows = []
        frame.mark(f"Directory: {current}", frame.curses.A_BOLD)
        error = validation_error()
        if error is None:
            frame.mark("  This directory is a valid choice. Press Enter.",
                       frame.curses.A_DIM)
        else:
            frame.mark(f"  {error}", frame.curses.A_BOLD)
        frame.mark("")
        if not entries:
            frame.mark("  (no subdirectories)")
        for entry in entries:
            frame.mark(f"  {entry.name}/")
        header = 3  # directory line, validity line, blank separator
        frame.cursor = header + (cursor if entries else 0)
        frame.message = ""
        frame.draw()
        curses = frame.curses
        key = frame.get_key(cancel)
        if key in (curses.KEY_UP, ord("k")):
            cursor = max(0, cursor - 1)
        elif key in (curses.KEY_DOWN, ord("j")):
            if entries:
                cursor = min(len(entries) - 1, cursor + 1)
        elif key in (curses.KEY_RIGHT, ord("l")):
            if entries:
                current = entries[cursor]
                cursor = 0
        elif key in (curses.KEY_LEFT, ord("h"), ord("u"),
                     curses.KEY_BACKSPACE, 8, 127):
            parent = current.parent
            if parent != current:
                current = parent
                cursor = 0
        elif key == ord("e"):
            result = frame.edit_line("", prompt="path: ")
            if result is not None:
                candidate = Path(os.path.expanduser(result))
                if not candidate.is_absolute():
                    candidate = current / candidate
                try:
                    candidate = candidate.resolve()
                except OSError:
                    pass
                if candidate.is_dir():
                    current = candidate
                    cursor = 0
                else:
                    frame.message = f"Not a directory: {candidate}"
                    frame.draw()
                    frame.get_key(cancel)
                    frame.get_key(cancel)
        elif key in (10, 13):  # Enter: accept the listed directory
            error = validation_error()
            if error is None:
                return current
            frame.message = f"{error}  (keep browsing)"
            frame.draw()
            frame.get_key(cancel)


# ---------------------------------------------------------------------------
# Widget: expandable checkbox tree
# ---------------------------------------------------------------------------

def checkbox_tree(scr, title: str, families: List[dict],
                  footer: Optional[str] = None,
                  expand_all: bool = False) -> List[Tuple[int, str]]:
    """Pick model families and packages from an expandable tree.

    FAMILIES is a list of dicts (one per family) shaped like::

        {
            "label": "Qwen3-TTS (qwen3_tts)",
            "detail": "tts, cloning, design",
            "options": [
                {"key": "Base-GGUF", "label": "base", "recommended": True},
                {"key": "VoiceDesign-GGUF", "label": "voicedesign",
                 "recommended": False},
            ],
        }

    Space on a family row checks its recommended option (or clears every
    option when one is already checked); Space on an option row toggles
    that option. Tab/Right expands or collapses the family under the
    cursor. Enter returns the flat list of (family_index, option_key)
    pairs for every checked option, in tree order; at least one checked
    option is required. The first family's recommended option starts
    checked (the prompt flow's default), and with EXPAND_ALL every
    family starts expanded.
    """
    footer = footer or ("Up/Down = move  Tab/Right = expand  Space = check  "
                        "Enter = accept  Esc = cancel")
    frame = Frame(scr, title, footer)
    cancel = (27, ord("q"))
    expanded = {index for index in range(len(families))} if expand_all else set()
    checked = set()  # (family_index, option_key)

    if families:
        expanded.add(0)
        first = families[0]["options"]
        for option in first:
            if option.get("recommended"):
                checked.add((0, option["key"]))
                break
        else:
            if first:
                checked.add((0, first[0]["key"]))

    def family_checked(index: int) -> bool:
        return any(pair[0] == index for pair in checked)

    def accept() -> List[Tuple[int, str]]:
        return [(index, option["key"])
                for index, family in enumerate(families)
                for option in family["options"]
                if (index, option["key"]) in checked]

    def visible_nodes() -> List[tuple]:
        nodes: List[tuple] = []  # ("family", i) or ("option", i, key)
        for index, family in enumerate(families):
            nodes.append(("family", index))
            if index in expanded:
                for option in family["options"]:
                    nodes.append(("option", index, option["key"]))
        return nodes

    cursor = 0
    while True:
        nodes = visible_nodes()
        cursor = max(0, min(cursor, len(nodes) - 1))
        frame.rows = []
        for node in nodes:
            if node[0] == "family":
                index = node[1]
                family = families[index]
                mark = "x" if family_checked(index) else " "
                arrow = "-" if index in expanded else "+"
                attr = frame.curses.A_BOLD if family_checked(index) else 0
                frame.mark(f"[{mark}] {arrow} {family['label']}", attr)
            else:
                _, index, option_key = node
                option = next(opt for opt in families[index]["options"]
                              if opt["key"] == option_key)
                is_on = (index, option_key) in checked
                mark = "x" if is_on else " "
                note = "  [recommended]" if option.get("recommended") else ""
                frame.mark(f"    [{mark}] {option['label']}{note}")
        frame.cursor = cursor
        node = nodes[cursor]
        frame.message = families[node[1]].get("detail", "")
        frame.message_attr = frame.curses.A_DIM
        frame.draw()
        curses = frame.curses
        key = frame.get_key(cancel)
        if key in (curses.KEY_UP, ord("k")):
            cursor = (cursor - 1) % len(nodes)
        elif key in (curses.KEY_DOWN, ord("j")):
            cursor = (cursor + 1) % len(nodes)
        elif key in (9, curses.KEY_RIGHT, ord("l")) and node[0] == "family":
            index = node[1]
            if index in expanded:
                expanded.discard(index)
            else:
                expanded.add(index)
        elif key == curses.KEY_LEFT and node[0] == "family":
            expanded.discard(node[1])
        elif key == ord(" "):
            if node[0] == "family":
                index = node[1]
                options = families[index]["options"]
                if family_checked(index):
                    for option in options:
                        checked.discard((index, option["key"]))
                else:
                    for option in options:
                        if option.get("recommended"):
                            checked.add((index, option["key"]))
                            break
                    else:
                        if options:
                            checked.add((index, options[0]["key"]))
                    expanded.add(index)
            else:
                _, index, option_key = node
                if (index, option_key) in checked:
                    checked.discard((index, option_key))
                else:
                    checked.add((index, option_key))
        elif key in (10, 13):  # Enter: accept the checked selection
            selection = accept()
            if selection:
                return selection
            frame.message = "Check at least one model package (Space)"
            frame.message_attr = None
            frame.draw()
            frame.get_key(cancel)
            frame.message_attr = frame.curses.A_DIM