Pyra/src/pyra/setup/wizard.py

import httpx
import questionary
from rich.console import Console
from rich.panel import Panel
from rich.text import Text

from pyra.config.manager import save_config
from pyra.config.schema import GeneralConfig, ProviderConfig, PyraConfig
from pyra.setup.providers import PROVIDERS, Provider, get_provider

console = Console()

_USE_CASE_PLUGINS: dict[str, list[str]] = {
    "Research & web":           ["websearch", "headless_browser"],
    "Development & servers":    ["server_manager", "ssh_tool", "docker_tool"],
    "File management":          ["gdrive", "onedrive", "dropbox_tool"],
    "Communication bots":       ["matrix_bot", "telegram_bot", "signal_bot"],
    "Email":                    ["email"],
    "Productivity & calendars": ["nextcloud"],
}


def run_setup() -> None:
    console.print(Panel(
        Text("Welcome to Pyra Setup", justify="center", style="bold cyan"),
        subtitle="Personal AI Assistant",
        border_style="cyan",
    ))
    console.print()

    user_name, purpose, use_cases = _collect_user_profile()

    provider = _choose_provider()
    model = _choose_model(provider)

    if provider.requires_key:
        _collect_api_key(provider)

    _test_connection(provider, model)

    cfg = PyraConfig(
        ai=ProviderConfig(
            provider_id=provider.id,
            model=model,
            base_url=provider.base_url,
        ),
        general=GeneralConfig(user_name=user_name, purpose=purpose),
    )
    save_config(cfg)

    _suggest_plugins(use_cases)

    console.print()
    console.print(Panel(
        f"[green]Setup complete![/green]\n\n"
        f"Provider: [bold]{provider.display_name}[/bold]\n"
        f"Model:    [bold]{model}[/bold]\n\n"
        "Run [bold cyan]pyra chat[/bold cyan] to start talking.",
        border_style="green",
    ))


def _collect_user_profile() -> tuple[str, str, list[str]]:
    console.print("[bold]Let's personalise your setup.[/bold]")
    console.print()

    name = questionary.text("What should Pyra call you?", default="User").ask()
    if name is None:
        raise SystemExit(0)
    name = name.strip() or "User"

    purpose = questionary.text(
        "In one sentence, what will you mainly use Pyra for? (optional)",
    ).ask()
    if purpose is None:
        raise SystemExit(0)
    purpose = purpose.strip()

    use_cases = questionary.checkbox(
        "Which areas interest you? (Space to select, Enter to confirm)",
        choices=list(_USE_CASE_PLUGINS.keys()),
    ).ask()
    if use_cases is None:
        raise SystemExit(0)

    console.print()
    return name, purpose, use_cases or []


def _suggest_plugins(use_cases: list[str]) -> None:
    if not use_cases:
        return

    lines: list[str] = []
    for uc in use_cases:
        plugins = _USE_CASE_PLUGINS.get(uc, [])
        if plugins:
            lines.append(f"[bold]{uc}[/bold]")
            for p in plugins:
                lines.append(f"  pyra plugin install {p}")
    if not lines:
        return

    lines.append("")
    lines.append("[dim]All listed plugins are in development — install when available.[/dim]")

    console.print()
    console.print(Panel(
        "\n".join(lines),
        title="Suggested plugins",
        border_style="dim cyan",
    ))


def _choose_provider() -> Provider:
    local = [p for p in PROVIDERS if p.group == "Local"]
    cloud = [p for p in PROVIDERS if p.group == "Cloud"]

    choices = (
        [questionary.Choice("── Local ──────────────────", disabled=True)]
        + [questionary.Choice(p.display_name, value=p.id) for p in local]
        + [questionary.Choice("── Cloud ──────────────────", disabled=True)]
        + [questionary.Choice(p.display_name, value=p.id) for p in cloud]
    )

    provider_id = questionary.select(
        "Choose your AI provider:",
        choices=choices,
    ).ask()

    if provider_id is None:
        raise SystemExit(0)

    provider = get_provider(provider_id)

    if provider.connectivity_check:
        _check_local_server(provider)

    return provider


def _check_local_server(provider: Provider) -> None:
    console.print(f"  Checking connection to [bold]{provider.display_name}[/bold]...", end=" ")
    try:
        resp = httpx.get(provider.connectivity_check, timeout=3.0)
        resp.raise_for_status()
        console.print("[green]✓[/green]")
    except Exception:
        console.print("[yellow]✗ (server not reachable)[/yellow]")
        console.print(
            f"  [yellow]Warning:[/yellow] Could not reach {provider.base_url}.\n"
            f"  Make sure {provider.display_name} is running before using Pyra."
        )


def _choose_model(provider: Provider) -> str:
    model = questionary.text(
        "Model name:",
        default=provider.default_model,
    ).ask()
    if model is None:
        raise SystemExit(0)
    return model.strip()


def _collect_api_key(provider: Provider) -> None:
    from pyra.vault.writer import set_key

    console.print(
        f"\n  [dim]API key will be stored in the encrypted vault — never in config.yaml[/dim]"
    )
    key = questionary.password(f"Enter your {provider.display_name} API key:").ask()
    if key is None:
        raise SystemExit(0)
    key = key.strip()
    if not key:
        console.print("[red]No key entered — skipping.[/red]")
        return
    set_key(provider.id, key)
    console.print("  [green]✓ Key stored in vault[/green]")


def _test_connection(provider: Provider, model: str) -> None:
    from pyra.vault.reader import get_key

    console.print("\n  Running test call...", end=" ")
    try:
        import litellm

        # Local providers don't need a real key but litellm still requires the field
        api_key = get_key(provider.id) if provider.requires_key else "local"
        kwargs: dict = {
            "model": f"{provider.litellm_prefix}{model}",
            "messages": [{"role": "user", "content": "Reply with exactly: OK"}],
            "max_tokens": 10,
            "api_key": api_key,
        }
        if provider.base_url:
            kwargs["api_base"] = provider.base_url

        litellm.completion(**kwargs)
        console.print("[green]✓ Connection OK[/green]")
    except Exception as exc:
        console.print(f"[yellow]✗ Test call failed: {exc}[/yellow]")
        console.print("  [dim]You can still proceed — check your config with 'pyra setup' again.[/dim]")