dc-nogi/bot/display.py

import asyncio
import time
from typing import List, Optional

from rich.console import Console, Group
from rich.live import Live
from rich.panel import Panel
from rich.progress_bar import ProgressBar
from rich.spinner import Spinner
from rich.table import Table
from rich.text import Text

from bot.storage import Storage

_TYPE_LABELS = {"guild": "server", "group_dm": "group", "dm": "DM"}


class CountdownBar:
    """A Rich renderable that computes its own remaining time at render
    time (same trick Spinner uses to animate for free under a Live display)
    so it ticks down smoothly without us needing an explicit polling loop."""

    def __init__(self, duration: float, label: str):
        self.duration = max(duration, 0.001)
        self.label = label
        self.start_time = time.monotonic()

    def __rich_console__(self, console, options):
        elapsed = time.monotonic() - self.start_time
        remaining = max(0.0, self.duration - elapsed)
        bar = ProgressBar(total=self.duration, completed=min(elapsed, self.duration), width=30)
        grid = Table.grid(padding=(0, 1))
        grid.add_column()
        grid.add_column(ratio=1)
        grid.add_column(justify="right")
        grid.add_row(Text(f"⏳ {self.label}", style="dim"), bar, Text(f"{remaining:0.1f}s", style="dim"))
        yield grid


def _render(
    stats: List[dict],
    friend_label: Optional[str] = None,
    status: Optional[str] = None,
    countdown: Optional[CountdownBar] = None,
    active_spinner: Optional[Spinner] = None,
    active_target_id: Optional[str] = None,
) -> Panel:
    table = Table(expand=True)
    table.add_column("Target")
    table.add_column("Type")
    table.add_column("Progress", ratio=1)
    table.add_column("Reacted/Found", justify="right")
    table.add_column("Last reaction (UTC)")

    total_reacted = 0
    total_found = 0
    for row in stats:
        reacted = row["reaction_count"]
        found = max(row["messages_found"], reacted)  # found can never be less than reacted
        total_reacted += reacted
        total_found += found

        name = row["name"] or row["target_id"]
        if active_spinner is not None and active_target_id is not None and str(row["target_id"]) == str(active_target_id):
            target_cell = _spinner_row(active_spinner, name)
        else:
            target_cell = Text(name)

        bar = ProgressBar(total=max(found, reacted, 1), completed=reacted, width=None)
        table.add_row(
            target_cell,
            _TYPE_LABELS.get(row["target_type"], row["target_type"]),
            bar,
            f"{reacted}/{found}",
            row["last_reaction_at"] or "-",
        )

    if not stats:
        table.add_row("(no targets discovered yet)", "-", "-", "-", "-")

    # Status/countdown rows are always present (even blank) so the panel's
    # height never changes as they toggle on/off — that jumping is jarring
    # under a Live display that also has scrolling log output above it.
    status_line = Text(f"▸ {status}", style="dim italic") if status else Text(" ")
    countdown_line = countdown if countdown is not None else Text(" ")
    content = Group(table, status_line, countdown_line)

    who = f" for {friend_label}" if friend_label else ""
    return Panel(content, title=f"Reaction stats{who} — total: {total_reacted}/{total_found}", border_style="cyan")


def _spinner_row(spinner: Spinner, name: str):
    grid = Table.grid(padding=(0, 1))
    grid.add_column()
    grid.add_column()
    grid.add_row(spinner, Text(name))
    return grid


class StatsDisplay:
    """A stats panel pinned to the bottom of the terminal, with normal
    log/print output scrolling above it (via rich.live.Live)."""

    def __init__(self, console: Console, storage: Storage):
        self.console = console
        self.storage = storage
        self.friend_label: Optional[str] = None
        self.status: str = "Starting..."
        self.active_target_id: Optional[str] = None
        self._spinner = Spinner("dots")
        self._countdown: Optional[CountdownBar] = None
        self._last_stats: List[dict] = []
        self._live = Live(
            _render([], None, self.status),
            console=console,
            refresh_per_second=8,
            transient=False,
        )

    def __enter__(self):
        self._live.__enter__()
        return self

    def __exit__(self, *exc_info):
        return self._live.__exit__(*exc_info)

    def _redraw(self) -> None:
        self._live.update(
            _render(
                self._last_stats,
                self.friend_label,
                self.status,
                self._countdown,
                self._spinner,
                self.active_target_id,
            )
        )

    def set_friend(self, label: str) -> None:
        self.friend_label = label

    def set_active_target(self, target_id: Optional[object]) -> None:
        """Show a little spinner next to whichever target row is currently
        being scanned/reacted to. Pass None when idle (e.g. once backlog
        scanning finishes and we're just watching for live messages)."""
        self.active_target_id = str(target_id) if target_id is not None else None
        self._redraw()

    def set_status(self, text: str) -> None:
        """Update the "what's happening right now" line without hitting the
        DB — cheap enough to call for every jittered wait / API call."""
        self.status = text
        self._redraw()

    def start_countdown(self, duration: float, label: str) -> None:
        """Show a live-ticking countdown bar for a jittered wait, replacing
        the status line for the duration of the wait."""
        self._countdown = CountdownBar(duration, label)
        self._redraw()

    def clear_countdown(self) -> None:
        self._countdown = None
        self._redraw()

    def refresh_sync(self) -> None:
        self._last_stats = self.storage.get_stats()
        self._redraw()

    async def refresh(self) -> None:
        self._last_stats = await asyncio.to_thread(self.storage.get_stats)
        self._redraw()