takopi/scripts/onboarding_preview.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Iterable, Iterator

import anyio
from rich.console import Console
from rich.panel import Panel
from rich.text import Text

from takopi.config import ConfigError
from takopi.telegram import onboarding as ob
from takopi.telegram.api_models import User


def section(console: Console, title: str) -> None:
    console.print("")
    console.print(f"=== {title} ===", markup=False)


def render_confirm(console: Console, prompt: str) -> None:
    console.print(f"? {prompt} (yes/no)", markup=False)


def render_password(console: Console, prompt: str) -> None:
    console.print(f"? {prompt} {'*' * 28}", markup=False)


def render_select(console: Console, prompt: str, choices: list[str]) -> None:
    console.print(f"? {prompt} (use arrow keys)", markup=False)
    for index, choice in enumerate(choices):
        marker = ">" if index == 0 else " "
        console.print(f"{marker} {choice}", markup=False)


def next_value(values: Iterator[Any], label: str) -> Any:
    try:
        return next(values)
    except StopIteration as exc:
        raise RuntimeError(f"scripted ui ran out of {label} responses") from exc


class ScriptedUI:
    def __init__(
        self,
        console: Console,
        *,
        confirms: Iterable[bool | None],
        selects: Iterable[Any],
        passwords: Iterable[str | None],
    ) -> None:
        self._console = console
        self._confirms = iter(confirms)
        self._selects = iter(selects)
        self._passwords = iter(passwords)

    @property
    def console(self) -> Console:
        return self._console

    def panel(
        self,
        title: str | None,
        body: str,
        *,
        border_style: str = "yellow",
    ) -> None:
        panel = Panel(
            body,
            title=title,
            border_style=border_style,
            padding=(1, 2),
            expand=False,
        )
        self._console.print(panel)

    def step(self, title: str, *, number: int) -> None:
        self._console.print("")
        self._console.print(Text(f"step {number}: {title}", style="bold yellow"))
        self._console.print("")

    def print(self, text: object = "", *, markup: bool | None = None) -> None:
        if markup is None:
            self._console.print(text)
            return
        self._console.print(text, markup=markup)

    async def confirm(self, prompt: str, default: bool = True) -> bool | None:
        render_confirm(self._console, prompt)
        return next_value(self._confirms, "confirm")

    async def select(self, prompt: str, choices: list[tuple[str, Any]]) -> Any | None:
        rendered = [label for label, _value in choices]
        render_select(self._console, prompt, rendered)
        return next_value(self._selects, "select")

    async def password(self, prompt: str) -> str | None:
        render_password(self._console, prompt)
        return next_value(self._passwords, "password")


@dataclass
class ScriptedServices:
    bot: User
    chat: ob.ChatInfo
    engines: list[tuple[str, bool, str | None]]
    topics_issue: ConfigError | None = None
    existing_config: dict[str, Any] | None = None
    written_config: dict[str, Any] | None = None

    async def get_bot_info(self, _token: str) -> User | None:
        return self.bot

    async def wait_for_chat(self, _token: str) -> ob.ChatInfo:
        return self.chat

    async def validate_topics(
        self, _token: str, _chat_id: int, _scope: ob.TopicScope
    ) -> ConfigError | None:
        return self.topics_issue

    def list_engines(self) -> list[tuple[str, bool, str | None]]:
        return self.engines

    def read_config(self, _path) -> dict[str, Any]:
        return dict(self.existing_config or {})

    def write_config(self, _path, data: dict[str, Any]) -> None:
        self.written_config = data


async def run_flow(title: str, ui: ScriptedUI, svc: ScriptedServices) -> None:
    section(ui.console, title)
    state = ob.OnboardingState(config_path=ob.HOME_CONFIG_PATH, force=False)
    await ob.run_onboarding(ui, svc, state)


def main() -> None:
    console = Console()

    bot = User(id=1, username="bunny_agent_bot", first_name="Bunny")
    group_chat = ob.ChatInfo(
        chat_id=-1001234567890,
        username=None,
        title="takopi devs",
        first_name=None,
        last_name=None,
        chat_type="supergroup",
    )
    private_chat = ob.ChatInfo(
        chat_id=462722,
        username="banteg",
        title=None,
        first_name="Banteg",
        last_name=None,
        chat_type="private",
    )
    engines_installed = [
        ("codex", True, "brew install codex"),
        ("claude", True, "brew install claude"),
        ("opencode", False, "brew install opencode"),
    ]
    engines_missing = [
        ("codex", False, "brew install codex"),
        ("claude", False, "brew install claude"),
        ("opencode", False, "brew install opencode"),
    ]

    anyio.run(
        run_flow,
        "assistant mode (private chat)",
        ScriptedUI(
            console,
            confirms=[True, True],
            selects=["assistant", "codex"],
            passwords=["123456789:ABCdef"],
        ),
        ScriptedServices(bot=bot, chat=private_chat, engines=engines_installed),
    )

    anyio.run(
        run_flow,
        "handoff mode (token instructions)",
        ScriptedUI(
            console,
            confirms=[False, True],
            selects=["handoff", "codex"],
            passwords=["123456789:ABCdef"],
        ),
        ScriptedServices(bot=bot, chat=private_chat, engines=engines_installed),
    )

    anyio.run(
        run_flow,
        "workspace mode (topics)",
        ScriptedUI(
            console,
            confirms=[True, True],
            selects=["workspace", "codex"],
            passwords=["123456789:ABCdef"],
        ),
        ScriptedServices(bot=bot, chat=group_chat, engines=engines_installed),
    )

    anyio.run(
        run_flow,
        "topics validation warning",
        ScriptedUI(
            console,
            confirms=[True, True],
            selects=["workspace", "assistant", "codex"],
            passwords=["123456789:ABCdef"],
        ),
        ScriptedServices(
            bot=bot,
            chat=group_chat,
            engines=engines_installed,
            topics_issue=ConfigError("bot is missing admin rights"),
        ),
    )

    anyio.run(
        run_flow,
        "no engines installed",
        ScriptedUI(
            console,
            confirms=[True, False],
            selects=["assistant"],
            passwords=["123456789:ABCdef"],
        ),
        ScriptedServices(bot=bot, chat=private_chat, engines=engines_missing),
    )

if __name__ == "__main__":
    main()