scrypt/ui/totem.py

"""Carving the totem: the first-run model download as a ritual.

A ~20GB download is the worst onboarding moment a game can have, so it
becomes a scene: the Warden's totem is carved row by row as the GGUF
arrives. The real work happens in a thread; the UI only watches a float.

Dismisses with True (carved) or False (cancelled / failed).
"""

from __future__ import annotations

import asyncio

from rich.text import Text
from textual import events
from textual.app import ComposeResult
from textual.containers import Vertical
from textual.screen import Screen
from textual.widgets import Static

from scrypt.inference import local
from scrypt.ui import palette as pal

TOTEM = [
    "      ▄▄▄█▄▄▄      ",
    "     ▟▀▀▀▀▀▀▀▙     ",
    "     █  ◉ ◉  █     ",
    "     █   ┃   █     ",
    "     █  ╰─╯  █     ",
    "      ▜▄▄▄▄▄▛      ",
    "     ▄█▀▀▀▀▀█▄     ",
    "     █ ▓▓▓▓▓ █     ",
    "     █ ▓ ⚖ ▓ █     ",
    "     █ ▓▓▓▓▓ █     ",
    "      ▀▀█▄█▀▀      ",
    "     ▄▄█▀▀▀█▄▄     ",
    "     █ ░░░░░ █     ",
    "      ▀▀▀▀▀▀▀      ",
]
RAW_WOOD = "▒" * 13


class CancelledCarving(Exception):
    pass


class TotemScreen(Screen):
    CSS = """
    #totem-title { height: 3; content-align: center middle; }
    #totem { height: 16; content-align: center middle; }
    #totem-bar { height: 2; content-align: center middle; }
    #totem-prompt { height: 1; dock: bottom; background: $panel; content-align: center middle; }
    """

    def __init__(self) -> None:
        super().__init__()
        self._frac = 0.0
        self._cancel = False
        self._error: str | None = None
        self._last_carved = 0
        self.quant = local.choose_quant() or "Q3_K_S"
        self.size_gb = next(
            (s for q, s, _ in local.QUANT_LADDER if q == self.quant), 20
        )

    def compose(self) -> ComposeResult:
        with Vertical():
            yield Static(id="totem-title")
            yield Static(id="totem")
            yield Static(id="totem-bar")
            yield Static(id="totem-prompt")

    def on_mount(self) -> None:
        self.query_one("#totem-title", Static).update(Text(
            "THE TOTEM MUST BE CARVED\n"
            f"the Warden's mind, {self.quant}, ~{self.size_gb}GB of it",
            style=f"bold {pal.WARDEN}", justify="center",
        ))
        self.query_one("#totem-prompt", Static).update(
            Text("[esc] stop carving — the game plays scripted until the totem is whole")
        )
        self.set_interval(0.2, self._redraw)
        self.run_worker(self._carve())
        self._redraw()

    # ------------------------------------------------------------- carve

    def _progress(self, frac: float) -> None:
        # Called from the download thread. A float store is all it does.
        if self._cancel:
            raise CancelledCarving()
        self._frac = frac

    async def _carve(self) -> None:
        try:
            await asyncio.to_thread(local.download_model, self._progress)
        except CancelledCarving:
            self.dismiss(False)
            return
        except Exception as err:
            self._error = str(err) or err.__class__.__name__
            self.query_one("#totem-title", Static).update(Text(
                f"the chisel slipped\n{self._error[:90]}",
                style=f"bold {pal.DANGER}", justify="center",
            ))
            self.query_one("#totem-prompt", Static).update(Text("[any key] retreat"))
            return
        self._frac = 1.0
        self._redraw()
        self.query_one("#totem-prompt", Static).update(
            Text("[any key] wake what you carved")
        )

    # ------------------------------------------------------------ redraw

    def _redraw(self) -> None:
        carved = int(self._frac * len(TOTEM))
        fresh_row = carved - 1 if carved > self._last_carved else None
        self._last_carved = max(self._last_carved, carved)
        art = Text(justify="center")
        for i, row in enumerate(TOTEM):
            if i < carved:
                if i == fresh_row:  # chips fly off the newest cut
                    art.append("✦" + row[1:-1] + "✦\n", style=pal.GOLD)
                else:
                    art.append(row + "\n", style=pal.JADE_GLOW)
            else:
                art.append(f"   {RAW_WOOD}   \n", style=pal.GHOST)
        self.query_one("#totem", Static).update(art)
        width = 30
        filled = int(self._frac * width)
        bar = Text(justify="center")
        bar.append("█" * filled, style=pal.GREEN)
        bar.append("▒" * (width - filled), style=pal.GHOST)
        bar.append(f"  {self._frac * 100:3.0f}%", style=pal.MUTED)
        self.query_one("#totem-bar", Static).update(bar)

    # -------------------------------------------------------------- keys

    def on_key(self, event: events.Key) -> None:
        if self._frac >= 1.0 and self._error is None:
            self.dismiss(True)
        elif self._error is not None:
            self.dismiss(False)
        elif event.key == "escape":
            self._cancel = True  # the download thread raises on next chunk