afterglow/ui/menus.py

Part of Afterglow.

"""Afterglow menu screens.

Standalone, responsive ``Control`` screens built on the ``simvx.core.ui``
toolkit. Each screen is self-contained: it owns its layout, draws its own dark
on-theme panel + glowing accents, and exposes plain-callable hooks the game sets
(``screen.on_play = handler``). No screen reaches into game state directly; the
game wires the hooks and feeds in a ``Progress`` for state-driven screens (level
select, options, results).

All screens are anchored (``anchor=``) so they scale to any aspect
ratio, keyboard / controller navigable (buttons take focus and are linked in tab
order), and touch-tappable (buttons respond to ``Mouse.LEFT``, which the
web runtime surfaces touch through).

Public screens:
  * :class:`TitleScreen`   -- wordmark + Play / Levels / Options / Quit.
  * :class:`LevelSelect`   -- the three worlds, rooms with tick / time / shard.
  * :class:`PauseMenu`     -- Resume / Restart / Options / Quit to map.
  * :class:`OptionsMenu`   -- volume sliders + toggles + assist group.
  * :class:`ResultsScreen` -- per-world / final completion summary.
"""

from __future__ import annotations

from collections.abc import Callable, Sequence

from simvx.core import (
    Anchor,
    Button,
    CheckBox,
    Control,
    FocusMode,
    HBoxContainer,
    Key,
    KeyEvent,
    Label,
    Panel,
    PointerFilter,
    Slider,
    VBoxContainer,
)
from simvx.core.ui.navigation import UI_ACCEPT, ui_action_matches

from .hud import format_time

#: Optional sound sink. The game assigns a callable taking an SFX name
#: (``"menu_move"`` when focus walks, ``"menu_confirm"`` on an activation) so the
#: menus are audible; left unset they are silent, which keeps this module free of
#: any audio dependency for tests and reuse.
play_sound: Callable[[str], None] | None = None


def _sound(name: str) -> None:
    if play_sound is not None:
        play_sound(name)


# Shared palette: dark panels, glowing accent. Kept here (not a Theme) so this
# module stays a single self-contained file.
ACCENT = (0.55, 1.0, 0.78, 1.0)
ACCENT_DIM = (0.40, 0.72, 0.58, 1.0)
GOLD = (1.0, 0.78, 0.30, 1.0)
TEXT = (0.93, 0.97, 0.95, 1.0)
TEXT_DIM = (0.62, 0.68, 0.70, 1.0)
LOCKED = (0.45, 0.47, 0.52, 1.0)
PANEL_BG = (0.05, 0.06, 0.09, 0.92)
ROW_BG = (0.09, 0.11, 0.15, 0.9)
BACKDROP = (0.02, 0.03, 0.05, 0.78)


#: Keys that move focus down / up the chain or back out of the screen. Tab
#: pairs with down and Shift+Tab with up, matching the engine's own tab order,
#: which is why the two are resolved by :func:`_moves_next` rather than listed.
#: Activation is not here either: it is the engine's ``ui_accept`` action,
#: asked for per event (see :func:`_is_activate`).
_NEXT_KEYS = {Key.DOWN, Key.RIGHT}
_PREV_KEYS = {Key.UP, Key.LEFT}
#: Split directions for 2D (grid) navigation. A button with ``focus_left`` /
#: ``focus_right`` set uses them for left/right; otherwise left/right fall back to
#: the vertical chain (so plain vertical menus are unaffected).
_DOWN_KEYS = {Key.DOWN}
_UP_KEYS = {Key.UP}
_RIGHT_KEYS = {Key.RIGHT}
_LEFT_KEYS = {Key.LEFT}
_BACK_KEYS = {Key.ESCAPE}
#: Horizontal keys a focused slider consumes to nudge its value (instead of
#: walking the focus chain). Buttons / checkboxes keep using them to move focus.
_ADJUST_KEYS = {Key.LEFT, Key.RIGHT}
#: Fixed keys these controls act on. Claimed on the press as well as the release,
#: because the engine's own Tab traversal runs on the press: without the claim it
#: would move focus out from under the release we are waiting for, and the move
#: would happen silently. Escape is in the set for the same reason -- these screens
#: are plain scenes with their own back handler, and the press has to survive to
#: the release that acts on it.
_FIXED_OWNED_KEYS = (
    _NEXT_KEYS | _PREV_KEYS | _DOWN_KEYS | _UP_KEYS | _RIGHT_KEYS | _LEFT_KEYS | _BACK_KEYS | _ADJUST_KEYS
) | {Key.TAB}


def _moves_next(event: KeyEvent) -> bool:
    """True when this key walks the focus chain forwards. Tab does; Shift+Tab does not."""
    return not event.shift if event.key is Key.TAB else event.key in _NEXT_KEYS


def _moves_prev(event: KeyEvent) -> bool:
    """True when this key walks the focus chain backwards."""
    return event.shift if event.key is Key.TAB else event.key in _PREV_KEYS


def _moves_down(event: KeyEvent) -> bool:
    """The vertical half of :func:`_moves_next`, for a grid menu."""
    return not event.shift if event.key is Key.TAB else event.key in _DOWN_KEYS


def _moves_up(event: KeyEvent) -> bool:
    """The vertical half of :func:`_moves_prev`, for a grid menu."""
    return event.shift if event.key is Key.TAB else event.key in _UP_KEYS


def _is_activate(control: Control, event: KeyEvent) -> bool:
    """True when this key is bound to the engine's ``ui_accept`` action.

    Asked of the engine rather than compared against a literal, so a game that
    rebinds ``ui_accept`` moves menu confirmation with it, exactly as it moves a
    plain ``Button``'s. Space is therefore only ever a confirm key if the project
    binds it: the default is Enter, and Space stays the gameplay JUMP key, which
    is what stopped a screen opening under a held jump from confirming itself.
    """
    tree = control.tree
    return ui_action_matches(UI_ACCEPT, event, tree.input_map if tree is not None else None)


def _owns_key(control: Control, event: KeyEvent) -> bool:
    """True when this control acts on this key and so must claim it on the press."""
    return event.key in _FIXED_OWNED_KEYS or _is_activate(control, event)


class _MenuButton(Button):
    """A menu button that is keyboard- and controller-navigable when focused.

    The engine's base ``Button`` only reacts to ``Mouse.LEFT``, so the
    focused menu button drives its own directional navigation: up/down (and
    left/right/Tab/Shift+Tab) walk the ``focus_next``/``focus_previous`` chain
    ``_link_focus`` wired, and Escape/controller-B invokes the screen's back handler
    if one is set. Mouse / touch keep working through the base handler.

    Confirmation is the ENGINE's: the key is whatever ``ui_accept`` is bound to,
    and the activation itself is ``Button.activate()``. These screens still claim
    it on the press and act on the release, because that release-side bookkeeping
    is what makes the menu-opening key's orphan release a no-op (see below); the
    router's own press-side activation stands down for a claimed key, so the two
    never both fire.
    """

    on_back: Callable[[], None] | None = None

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # Keys whose PRESS this button observed while focused + in-tree. We only
        # act on a release if we saw its matching press: otherwise the very key
        # that OPENED this menu (an Escape press routed to the game, which spawns
        # the menu, then its release lands on the freshly-focused button) would
        # immediately fire "back" and the menu would flicker open->shut. Tracking
        # press ownership makes that orphan release a no-op.
        self._key_armed: set[Key] = set()
        # One confirm chime per activation, however it was triggered (mouse, touch,
        # Enter or controller-A all end up emitting ``pressed``).
        self.pressed.connect(self._play_confirm)

    @staticmethod
    def _play_confirm() -> None:
        _sound("menu_confirm")

    def on_gui_input(self, event):
        # Pointer and wheel: defer to the base button behaviour.
        if not isinstance(event, KeyEvent):
            return super().on_gui_input(event)
        if not self.focused:
            return None
        key = event.key
        if event.pressed:
            # Arm: remember we saw this key's press while focused, so its later
            # release is ours to act on (and not the menu-opening key's release).
            self._key_armed.add(key)
            # Act on key release, matching the engine's widget convention, but
            # claim the press now so nothing behind acts on it either.
            return _owns_key(self, event)
        if key not in self._key_armed:
            return None  # release with no matching press we owned: ignore (orphan)
        self._key_armed.discard(key)
        if _moves_down(event):
            self.focus_next_control()
            _sound("menu_move")
            return True
        elif _moves_up(event):
            self.focus_previous_control()
            _sound("menu_move")
            return True
        elif key in _RIGHT_KEYS:
            target = getattr(self, "focus_right", None)
            target.grab_focus() if target is not None else self.focus_next_control()
            _sound("menu_move")
            return True
        elif key in _LEFT_KEYS:
            target = getattr(self, "focus_left", None)
            target.grab_focus() if target is not None else self.focus_previous_control()
            _sound("menu_move")
            return True
        elif _is_activate(self, event):
            self.activate()
            return True
        elif key in _BACK_KEYS and self.on_back is not None:
            self.on_back()
            return True


class _KeyNav:
    """Mixin: arm-on-press / act-on-release keyboard navigation for a focused
    control whose base widget handles only the mouse (``Slider``, ``CheckBox``).

    The engine's non-modal key path delivers events solely to the focused
    control, and the base ``Slider`` / ``CheckBox`` ignore keys entirely, so a
    keyboard-only player who lands on one is stranded: arrows can't move focus and
    Escape can't back out (the Options menu trap). This mixin gives any such
    control the same up/down/Tab/Shift+Tab focus walk, Escape-to-back, and (for a slider)
    left/right adjust that :class:`_MenuButton` provides, using the identical
    arm-on-press bookkeeping so the menu-opening key's orphan release is ignored.

    Subclasses set ``on_back`` (Escape handler) and override ``_activate`` /
    ``_adjust`` for the ``ui_accept`` key and the horizontal (left/right) keys.
    """

    on_back: Callable[[], None] | None = None
    #: When True, left/right adjust the value instead of walking the focus chain
    #: (sliders); when False, left/right move focus like a button (checkboxes).
    _consumes_horizontal: bool = False

    def _init_keynav(self) -> None:
        self._key_armed: set[Key] = set()

    def _keynav_decided(self, event) -> bool | None:
        """What this mixin decided about a key event, for the widget to return.

        ``None`` means it did not take the event and the base widget should see
        it. ``True`` claims it, ``False`` stops it here without claiming: the
        mixin dealt with it, but nothing behind should be prevented from acting
        on a key it merely armed or ignored.
        """
        if not isinstance(event, KeyEvent):
            return None
        if self.tree is None or not self.focused:
            return False  # drop keys aimed at an orphaned / unfocused control
        key = event.key
        if event.pressed:
            self._key_armed.add(key)
            return _owns_key(self, event)
        if key not in self._key_armed:
            return False  # orphan release (e.g. the key that opened this menu)
        self._key_armed.discard(key)
        adjust = self._consumes_horizontal and key in _ADJUST_KEYS
        if adjust:
            self._adjust(+1 if key is Key.RIGHT else -1)
        elif _moves_next(event):
            self.focus_next_control()
            _sound("menu_move")
        elif _moves_prev(event):
            self.focus_previous_control()
            _sound("menu_move")
        elif _is_activate(self, event):
            self._activate()
        elif key in _BACK_KEYS and self.on_back is not None:
            self.on_back()
        else:
            return False
        return True

    # Subclass hooks (no-ops by default).
    def _activate(self) -> None: ...

    def _adjust(self, direction: int) -> None: ...


class _MenuSlider(_KeyNav, Slider):
    """A :class:`Slider` that is keyboard-navigable when focused: up/down/Tab/Shift+Tab
    walk the focus chain, left/right nudge the value by ``step``, Escape backs out."""

    _consumes_horizontal = True

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._init_keynav()

    def on_gui_input(self, event):
        decided = self._keynav_decided(event)
        if decided is not None:
            return decided
        return super().on_gui_input(event)

    def _adjust(self, direction: int) -> None:
        step = self.step if self.step > 0 else 0.05
        new_value = max(self.min_value, min(self.max_value, self.value + direction * step))
        if new_value != self.value:
            self.value = new_value
            self.value_changed.emit(self.value)


class _MenuCheckBox(_KeyNav, CheckBox):
    """A :class:`CheckBox` that is keyboard-navigable when focused: up/down/Tab/Shift+Tab
    walk the focus chain, ``ui_accept`` toggles, Escape backs out."""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._init_keynav()

    def on_gui_input(self, event):
        decided = self._keynav_decided(event)
        if decided is not None:
            return decided
        return super().on_gui_input(event)

    def _activate(self) -> None:
        self.activate()  # the engine's own toggle, exactly what a click does
        _sound("menu_confirm")


def _accent_button(text: str, on_press: Callable[[], None] | None = None, *, width: float = 280.0) -> _MenuButton:
    """A focusable, keyboard-navigable, on-theme menu button sized to a width."""
    btn = _MenuButton(text, on_press=on_press)
    btn.font_size = 20.0
    btn.size = (width, 46.0)
    btn.size_flags_horizontal = 0  # SHRINK_BEGIN: keep our width in a VBox
    btn.focus_mode = FocusMode.ALL
    btn.text_colour = TEXT
    btn.bg_colour = (0.10, 0.13, 0.18, 0.95)
    btn.hover_colour = (0.16, 0.30, 0.26, 1.0)
    btn.pressed_colour = (0.20, 0.40, 0.34, 1.0)
    btn.border_colour = ACCENT_DIM
    return btn


def _link_focus(buttons: Sequence[Control], *, on_back: Callable[[], None] | None = None) -> None:
    """Wire a vertical focus chain so up/down + tab cycle the buttons.

    ``on_back`` (when given) is attached to every :class:`_MenuButton` so Escape /
    controller-B backs out of the screen from any focused button.
    """
    n = len(buttons)
    for i, b in enumerate(buttons):
        b.focus_next = buttons[(i + 1) % n]
        b.focus_previous = buttons[(i - 1) % n]
        if on_back is not None and isinstance(b, _MenuButton | _KeyNav):
            b.on_back = on_back


def _link_grid_focus(
    columns: Sequence[Sequence[Control]], back: Control, *, on_back: Callable[[], None] | None = None
) -> None:
    """Wire 2D focus for a column grid: up/down within a column, left/right across.

    ``columns`` is one list of buttons per column. ``back`` sits below the grid:
    Down from a column's last button (or any vertical wrap) reaches it; left/right
    jump to the same-row button in the adjacent column (clamped to its length).
    """
    cols = [list(c) for c in columns if c]
    for col in cols:
        n = len(col)
        for i, b in enumerate(col):
            b.focus_next = col[i + 1] if i + 1 < n else back
            b.focus_previous = col[i - 1] if i > 0 else back
    if cols:
        back.focus_next = cols[0][0]
        back.focus_previous = cols[0][-1]
    for ci, col in enumerate(cols):
        left_col = cols[ci - 1] if ci > 0 else None
        right_col = cols[ci + 1] if ci + 1 < len(cols) else None
        for r, b in enumerate(col):
            b.focus_left = left_col[min(r, len(left_col) - 1)] if left_col else None
            b.focus_right = right_col[min(r, len(right_col) - 1)] if right_col else None
    back.focus_left = None
    back.focus_right = None
    if on_back is not None:
        for b in [*[x for c in cols for x in c], back]:
            if isinstance(b, _MenuButton | _KeyNav):
                b.on_back = on_back


def _centred_panel(width: float, height: float, *, name: str = "Panel") -> Panel:
    """A centred dark panel anchored to the screen centre with a fixed size."""
    panel = Panel(name=name, anchor=Anchor.CENTER, size=(width, height))
    panel.bg_colour = PANEL_BG
    panel.border_colour = ACCENT_DIM
    return panel


class _Screen(Control):
    """Base for full-screen menus: a dimmed backdrop filling the viewport."""

    def __init__(self, **kwargs):
        super().__init__(anchor=Anchor.FULL_RECT, **kwargs)
        self._backdrop = self.add_child(Panel(name="Backdrop", anchor=Anchor.FULL_RECT))
        self._backdrop.bg_colour = BACKDROP
        self._backdrop.border_colour = (0, 0, 0, 0)

    def focus_first(self) -> None:
        """Grab focus on the first focusable control (call after entering tree)."""
        first = self._first_focusable_descendant()
        if first is not None:
            first.grab_focus()


# ==========================================================================
# TitleScreen
# ==========================================================================


class TitleScreen(_Screen):
    """Splash / main menu: wordmark, tagline, and the primary actions.

    ``Play`` always begins at the first room; ``Continue`` (shown only when
    ``can_continue`` is set, i.e. there is saved progress) resumes at the first
    unfinished room.

    Hooks: ``on_play``, ``on_continue``, ``on_levels``, ``on_options``, ``on_quit``.
    """

    def __init__(self, can_continue: bool = False, **kwargs):
        super().__init__(**kwargs)
        self.on_play: Callable[[], None] | None = None
        self.on_continue: Callable[[], None] | None = None
        self.on_levels: Callable[[], None] | None = None
        self.on_options: Callable[[], None] | None = None
        self.on_quit: Callable[[], None] | None = None

        column = self.add_child(VBoxContainer(name="Column", anchor=Anchor.CENTER, size=(320, 420)))
        column.separation = 14.0
        column.alignment = "centre"

        wordmark = column.add_child(Label("AFTERGLOW", name="Wordmark"))
        wordmark.font_size = 56.0
        wordmark.alignment = "centre"
        wordmark.text_colour = ACCENT
        wordmark.size = (320, 70)

        tagline = column.add_child(Label("chase the light", name="Tagline"))
        tagline.font_size = 18.0
        tagline.alignment = "centre"
        tagline.text_colour = TEXT_DIM
        tagline.size = (320, 28)
        column.add_child(_spacer(18))

        # Returning players get a primary "Continue" (resume at the first unfinished
        # room); "Play" always starts a fresh run from the first room.
        buttons = []
        if can_continue:
            self.continue_button = _accent_button("Continue", self._fire("on_continue"))
            buttons.append(self.continue_button)
        self.play_button = _accent_button("Play", self._fire("on_play"))
        self.levels_button = _accent_button("Levels", self._fire("on_levels"))
        self.options_button = _accent_button("Options", self._fire("on_options"))
        self.quit_button = _accent_button("Quit", self._fire("on_quit"))
        buttons += [self.play_button, self.levels_button, self.options_button, self.quit_button]
        for b in buttons:
            column.add_child(b)
        _link_focus(buttons)

    def _fire(self, name: str) -> Callable[[], None]:
        return lambda: _call(getattr(self, name, None))


# ==========================================================================
# LevelSelect
# ==========================================================================


class LevelSelect(_Screen):
    """World / room picker rendered from a ``Progress`` and a ``worlds`` list.

    ``worlds`` is the ``rooms_data.WORLDS`` structure (a list of dicts with
    ``id``, ``name``, ``rooms``). Locked worlds render dimmed and non-clickable.

    Hooks: ``on_select(world_id, room_index)``, ``on_back``.
    """

    def __init__(self, worlds: Sequence[dict], progress=None, screen_size=None, **kwargs):
        super().__init__(**kwargs)
        self.on_select: Callable[[str, int], None] | None = None
        self.on_back: Callable[[], None] | None = None
        self._worlds = list(worlds)
        self._progress = progress
        self._screen_size = screen_size or (960.0, 600.0)
        self._room_buttons: list[Button] = []
        self.build()

    def set_progress(self, progress) -> None:
        """Swap the progress source and rebuild the list (call before showing)."""
        self._progress = progress
        self.build()

    def build(self) -> None:
        """Rebuild the picker as a per-world column layout sized to the viewport.

        The worlds sit side by side (a compact "world map"), so all rooms fit on
        screen without the picker running off the bottom. The panel is sized to the
        viewport, with a heading on top and a Back button below the columns.
        """
        for child in list(self.children):
            if child is not self._backdrop:
                self.remove_child(child)
        self._room_buttons = []

        sw, sh = float(self._screen_size[0]), float(self._screen_size[1])
        n_worlds = max(1, len(self._worlds))
        col_sep, inset, heading_h, back_h, sep = 16.0, 24.0, 34.0, 44.0, 12.0
        panel_w = min(360.0 + 260.0 * n_worlds, max(360.0, sw - 40.0))
        panel_h = min(600.0, max(300.0, sh - 40.0))
        inner_w = panel_w - 2 * inset
        col_w = (inner_w - col_sep * (n_worlds - 1)) / n_worlds
        cols_h = panel_h - 2 * inset - heading_h - back_h - 2 * sep

        panel = self.add_child(_centred_panel(panel_w, panel_h, name="LevelPanel"))
        column = panel.add_child(VBoxContainer(name="Column", anchor=Anchor.FULL_RECT, inset=inset))
        column.separation = sep

        heading = column.add_child(Label("SELECT A ROOM", name="Heading"))
        heading.font_size = 26.0
        heading.text_colour = ACCENT
        heading.size = (inner_w, heading_h)

        worlds_row = column.add_child(HBoxContainer(name="Worlds"))
        worlds_row.size = (inner_w, cols_h)
        worlds_row.separation = col_sep

        prog = self._progress
        grid: list[list[Control]] = []
        for world in self._worlds:
            wid = world["id"]
            unlocked = prog.is_world_unlocked(wid) if prog is not None else (wid == self._worlds[0]["id"])
            col = worlds_row.add_child(VBoxContainer(name=f"Col_{wid}"))
            col.size = (col_w, cols_h)
            col.separation = 6.0

            title = col.add_child(Label(world["name"], name=f"World_{wid}"))
            title.font_size = 18.0
            title.text_colour = TEXT if unlocked else LOCKED
            title.size = (col_w, 26)

            rooms = world.get("rooms", [])
            if not rooms:
                empty = col.add_child(Label("(coming soon)", name=f"Empty_{wid}"))
                empty.font_size = 13.0
                empty.text_colour = TEXT_DIM
                empty.size = (col_w, 22)
                continue
            if not unlocked:
                locked = col.add_child(Label("locked", name=f"Locked_{wid}"))
                locked.font_size = 13.0
                locked.text_colour = LOCKED
                locked.size = (col_w, 22)
                continue

            col_buttons: list[Control] = []
            for idx, room in enumerate(rooms):
                btn = self._room_row(wid, idx, room, prog, col_w)
                col.add_child(btn)
                col_buttons.append(btn)
            grid.append(col_buttons)

        self.back_button = _accent_button("Back", self._fire_back, width=200.0)
        self.back_button.size = (200.0, back_h)
        column.add_child(self.back_button)
        self._room_buttons = [*[b for c in grid for b in c], self.back_button]
        # Up/down walks a world column; left/right hops between worlds.
        _link_grid_focus(grid, self.back_button, on_back=self._fire_back)

    def _room_row(self, world_id: str, room_index: int, room, prog, width: float = 584.0) -> Control:
        """One room button labelled with its tick, best time, and shard dot."""
        done = prog.is_completed(world_id, room_index) if prog is not None else False
        shard = prog.has_shard(world_id, room_index) if prog is not None else False

        tick = "x" if done else " "
        dot = " *" if shard else ""
        name = getattr(room, "name", f"Room {room_index + 1}")
        label = f"[{tick}] {name}{dot}"

        btn = _accent_button(label, None, width=width)
        btn.font_size = 14.0
        btn.size = (width, 32.0)
        btn.alignment = "left"
        btn.text_colour = TEXT if done else TEXT_DIM
        btn.pressed.connect(lambda w=world_id, i=room_index: self._select(w, i))
        return btn

    def _select(self, world_id: str, room_index: int) -> None:
        if self.on_select is not None:
            self.on_select(world_id, room_index)

    def _fire_back(self) -> None:
        _call(self.on_back)


# ==========================================================================
# PauseMenu
# ==========================================================================


class PauseMenu(_Screen):
    """In-run pause overlay.

    Hooks: ``on_resume``, ``on_restart``, ``on_options``, ``on_quit``.
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.on_resume: Callable[[], None] | None = None
        self.on_restart: Callable[[], None] | None = None
        self.on_options: Callable[[], None] | None = None
        self.on_quit: Callable[[], None] | None = None

        panel = self.add_child(_centred_panel(340, 360, name="PausePanel"))
        column = panel.add_child(VBoxContainer(name="Column", anchor=Anchor.FULL_RECT, inset=(30, 28)))
        column.separation = 14.0
        column.alignment = "centre"

        heading = column.add_child(Label("PAUSED", name="Heading"))
        heading.font_size = 30.0
        heading.alignment = "centre"
        heading.text_colour = ACCENT
        heading.size = (280, 40)
        column.add_child(_spacer(8))

        self.resume_button = _accent_button("Resume", lambda: _call(self.on_resume))
        self.restart_button = _accent_button("Restart Room", lambda: _call(self.on_restart))
        self.options_button = _accent_button("Options", lambda: _call(self.on_options))
        self.quit_button = _accent_button("Quit to Map", lambda: _call(self.on_quit))
        buttons = [self.resume_button, self.restart_button, self.options_button, self.quit_button]
        for b in buttons:
            column.add_child(b)
        # Escape / controller-B resumes (closes the pause overlay) from any button.
        _link_focus(buttons, on_back=lambda: _call(self.on_resume))


# ==========================================================================
# OptionsMenu
# ==========================================================================


class OptionsMenu(_Screen):
    """Volume / accessibility / assist options, read & written via ``Progress``.

    All widgets read their initial value from ``progress.get_option`` and write
    back through ``progress.set_option`` on change, then fire ``on_change(key,
    value)`` so the game can apply the setting live (re-bus a volume, toggle
    fullscreen, enable an assist). Toggles whose option key is absent from the
    progress schema are skipped, so the screen degrades gracefully.

    Hooks: ``on_change(key, value)``, ``on_back``.
    """

    #: (option_key, label) for the simple boolean toggles.
    TOGGLES = (
        ("screenshake", "Screen shake"),
        ("photosensitive_safe", "Photosensitive-safe"),
    )
    #: (option_key, label) for the assist-mode group.
    ASSIST = (
        ("assist_mode", "Assist mode"),
        ("assist_slow_time", "Slow time"),
        ("assist_invincible", "Invincible"),
        ("assist_infinite_dash", "Infinite dash"),
        ("assist_extended_glow", "Extended glow"),
    )

    def __init__(self, progress=None, *, allow_fullscreen: bool = True, **kwargs):
        super().__init__(**kwargs)
        self.on_change: Callable[[str, object], None] | None = None
        self.on_back: Callable[[], None] | None = None
        self._progress = progress
        self._allow_fullscreen = allow_fullscreen
        self._focusables: list[Control] = []
        self.build()

    def set_progress(self, progress) -> None:
        """Swap the progress source and rebuild from its current option values."""
        self._progress = progress
        self.build()

    def build(self) -> None:
        for child in list(self.children):
            if child is not self._backdrop:
                self.remove_child(child)
        self._focusables = []

        panel = self.add_child(_centred_panel(520, 600, name="OptionsPanel"))
        column = panel.add_child(VBoxContainer(name="Column", anchor=Anchor.FULL_RECT, inset=(30, 24)))
        column.separation = 9.0

        heading = column.add_child(Label("OPTIONS", name="Heading"))
        heading.font_size = 28.0
        heading.text_colour = ACCENT
        heading.size = (460, 38)

        self._add_slider(column, "master_volume", "Master")
        self._add_slider(column, "sfx_volume", "SFX")
        self._add_slider(column, "music_volume", "Music")
        column.add_child(_spacer(6))

        if self._allow_fullscreen:
            self._add_toggle(column, "fullscreen", "Fullscreen", schema_optional=True)
        for key, label in self.TOGGLES:
            self._add_toggle(column, key, label)

        column.add_child(_spacer(6))
        assist_head = column.add_child(Label("Assist", name="AssistHead"))
        assist_head.font_size = 16.0
        assist_head.text_colour = GOLD
        assist_head.size = (460, 24)
        for key, label in self.ASSIST:
            self._add_toggle(column, key, label, schema_optional=True)

        column.add_child(_spacer(8))
        back = _accent_button("Back", lambda: _call(self.on_back), width=240.0)
        column.add_child(back)
        self._focusables.append(back)
        _link_focus(self._focusables, on_back=lambda: _call(self.on_back))

    def _opt(self, key: str, default):
        """Read an option, tolerating keys absent from the progress schema."""
        if self._progress is None:
            return default
        try:
            return self._progress.get_option(key)
        except KeyError:
            return default

    def _write(self, key: str, value) -> None:
        if self._progress is not None:
            try:
                self._progress.set_option(key, value)
            except KeyError:
                pass  # not in schema: still fire on_change so the game can apply
        if self.on_change is not None:
            self.on_change(key, value)

    def _add_slider(self, column, key: str, label: str) -> None:
        row = column.add_child(HBoxContainer(name=f"Row_{key}"))
        row.size = (460, 30)
        row.separation = 12.0
        cap = row.add_child(Label(label, name=f"Cap_{key}"))
        cap.font_size = 16.0
        cap.text_colour = TEXT
        cap.size = (130, 28)
        slider = row.add_child(_MenuSlider(0.0, 1.0, value=float(self._opt(key, 1.0))))
        slider.step = 0.05
        slider.size = (300, 22)
        slider.focus_mode = FocusMode.ALL
        slider.value_changed.connect(lambda v, k=key: self._write(k, float(v)))
        self._focusables.append(slider)

    def _add_toggle(self, column, key: str, label: str, *, schema_optional: bool = False) -> None:
        # Skip toggles whose key is not in the schema unless explicitly optional
        # (optional ones still fire on_change so the game may apply them live).
        if not schema_optional and self._progress is not None:
            try:
                self._progress.get_option(key)
            except KeyError:
                return
        checked = bool(self._opt(key, False))
        cb = column.add_child(_MenuCheckBox(label, checked=checked, name=f"Toggle_{key}"))
        cb.font_size = 16.0
        cb.text_colour = TEXT
        cb.size = (460, 26)
        cb.focus_mode = FocusMode.ALL
        cb.toggled.connect(lambda on, k=key: self._write(k, bool(on)))
        self._focusables.append(cb)


# ==========================================================================
# ResultsScreen
# ==========================================================================


class ResultsScreen(_Screen):
    """End-of-world / end-of-game completion summary.

    Call :meth:`set_summary` with the totals to render. Hooks: ``on_continue``,
    ``on_replay``, ``on_map``.
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.on_continue: Callable[[], None] | None = None
        self.on_replay: Callable[[], None] | None = None
        self.on_map: Callable[[], None] | None = None
        self._title = "WORLD COMPLETE"
        self._total_time = 0.0
        self._deaths = 0
        self._shards = 0
        self._shards_total = 0
        self._final = False
        self.build()

    def set_summary(
        self,
        *,
        title: str,
        total_time: float,
        deaths: int,
        shards: int,
        shards_total: int,
        final: bool = False,
    ) -> None:
        """Set the summary numbers and rebuild. ``final`` switches the wording."""
        self._title = title
        self._total_time = float(total_time)
        self._deaths = int(deaths)
        self._shards = int(shards)
        self._shards_total = int(shards_total)
        self._final = bool(final)
        self.build()

    def build(self) -> None:
        for child in list(self.children):
            if child is not self._backdrop:
                self.remove_child(child)

        panel = self.add_child(_centred_panel(460, 440, name="ResultsPanel"))
        column = panel.add_child(VBoxContainer(name="Column", anchor=Anchor.FULL_RECT, inset=(32, 28)))
        column.separation = 12.0
        column.alignment = "centre"

        heading = column.add_child(Label(self._title, name="Heading"))
        heading.font_size = 30.0
        heading.alignment = "centre"
        heading.text_colour = GOLD if self._final else ACCENT
        heading.size = (396, 40)
        column.add_child(_spacer(6))

        self._stat(column, "Total time", format_time(self._total_time), ACCENT)
        self._stat(column, "Deaths", str(self._deaths), TEXT)
        self._stat(column, "Shards", f"{self._shards} / {self._shards_total}", GOLD)
        column.add_child(_spacer(12))

        self.continue_button = _accent_button(
            "Next World" if not self._final else "Finish", lambda: _call(self.on_continue)
        )
        self.replay_button = _accent_button("Replay", lambda: _call(self.on_replay))
        self.map_button = _accent_button("Map", lambda: _call(self.on_map))
        buttons = [self.continue_button, self.replay_button, self.map_button]
        for b in buttons:
            column.add_child(b)
        _link_focus(buttons, on_back=lambda: _call(self.on_map))

    def _stat(self, column, label: str, value: str, value_colour) -> None:
        row = column.add_child(HBoxContainer(name=f"Stat_{label}"))
        row.size = (396, 28)
        cap = row.add_child(Label(label, name="Cap"))
        cap.font_size = 18.0
        cap.text_colour = TEXT_DIM
        cap.size = (220, 26)
        val = row.add_child(Label(value, name="Val"))
        val.font_size = 18.0
        val.alignment = "right"
        val.text_colour = value_colour
        val.size = (176, 26)


# -- shared helpers --------------------------------------------------------


def _spacer(height: float) -> Control:
    """An empty, non-interactive control used as vertical padding in a VBox."""
    s = Control(name="Spacer")
    s.size = (1.0, float(height))
    s.pointer_filter = PointerFilter.IGNORE
    return s


def _call(fn: Callable[[], None] | None) -> None:
    """Invoke a hook if one is set (no-op when the game left it unwired)."""
    if fn is not None:
        fn()


__all__ = [
    "TitleScreen",
    "LevelSelect",
    "PauseMenu",
    "OptionsMenu",
    "ResultsScreen",
    "ACCENT",
    "GOLD",
]