gpu-rent/src/gpu_rent/ux.py

"""User-facing prints: flavor list, cost warnings (no invented ₽)."""

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass
from typing import Any

from rich.markup import escape

from gpu_rent.config import Config
from gpu_rent.inventory import FlavorInfo, looks_like_gpu, rank_flavors, resolve_flavor

Log = Callable[[str], None]
Ask = Callable[[str, str], str]
Confirm = Callable[[str], bool]


def list_ranked_flavors(flavors: list[Any], cfg: Config) -> list[FlavorInfo]:
    gpu = [f for f in flavors if looks_like_gpu(f)]
    return rank_flavors(gpu or flavors, cfg.flavor_preference)


def format_flavor_lines(ranked: list[FlavorInfo], picked: FlavorInfo | None = None) -> list[str]:
    if not ranked:
        return ["[yellow]flavors: пусто (проверь регион / FLAVOR_PREFERENCE)[/yellow]"]
    lines = ["[bold cyan]flavors (по FLAVOR_PREFERENCE):[/bold cyan]"]
    for i, info in enumerate(ranked, 1):
        mark = " [green]← выберем[/green]" if picked and info.id == picked.id else ""
        ram = f"{info.ram_mb // 1024}GB" if info.ram_mb else "?"
        vcpu = str(info.vcpus) if info.vcpus is not None else "?"
        label = escape(info.label or "—")
        name = escape(info.name)
        fid = escape(info.id)
        lines.append(
            f"  [cyan]{i}.[/cyan] [bold]{label}[/bold] {name}  "
            f"vCPU={vcpu} RAM={ram}  id={fid}{mark}"
        )
    return lines


def cost_and_risk_lines(cfg: Config, *, spot: bool, flavor_name: str) -> list[str]:
    """Honest billing notes — OpenStack has no ₽ prices."""
    return [
        f"план: {'preemptible ' if spot else ''}{flavor_name}, data {cfg.data_volume_size_gb} GB",
        "₽: цены в OpenStack API нет — смотри панель Selectel (GPU ₽/час + диск ₽/мес).",
        "диск data тарифицируется 24/7, даже когда GPU выключен (stop).",
        f"idle-killer: {cfg.idle_grace_minutes} мин льготы после boot, потом "
        f"{cfg.idle_minutes} мин пустой очереди → delete compute. Отложить: gpu-rent hold",
        "preemptible: хостер может усыпить (~24 ч окно) → EXPIRED; tunnel сам unshelve, "
        "или gpu-rent up",
        "Ctrl+C на tunnel GPU не гасит — только gpu-rent stop или idle-killer. "
        "Опционально: gpu-rent watchdog install — аварийное закрытие окна/ребут "
        "после grace тоже stop (Ctrl+C по-прежнему detach)",
    ]


def print_up_preview(
    cfg: Config,
    flavors: list[Any],
    *,
    picked: FlavorInfo,
    spot: bool,
    log: Log,
) -> None:
    ranked = list_ranked_flavors(flavors, cfg)
    for line in format_flavor_lines(ranked, picked):
        log(line)
    log("")
    for line in cost_and_risk_lines(cfg, spot=spot, flavor_name=picked.name):
        log(f"! {line}")


@dataclass(frozen=True)
class ServerPlan:
    flavor: FlavorInfo
    data_gb: int
    spot: bool


def prompt_server_plan(
    cfg: Config,
    flavors: list[Any],
    *,
    picked: FlavorInfo,
    spot: bool,
    ask: Ask,
    confirm: Confirm | None = None,
) -> ServerPlan:
    """Interactive Selectel VM knobs before create (flavor / disk / spot)."""
    ranked = list_ranked_flavors(flavors, cfg)
    if not ranked:
        return ServerPlan(flavor=picked, data_gb=cfg.data_volume_size_gb, spot=spot)

    default_idx = 1
    for i, info in enumerate(ranked, 1):
        if info.id == picked.id:
            default_idx = i
            break

    raw_idx = ask(
        f"Flavor Selectel [1-{len(ranked)}] (Enter = рекомендация)",
        str(default_idx),
    ).strip()
    try:
        idx = int(raw_idx)
    except ValueError as exc:
        raise ValueError(f"номер flavor: жду 1…{len(ranked)}, получили {raw_idx!r}") from exc
    if idx < 1 or idx > len(ranked):
        raise ValueError(f"номер flavor: жду 1…{len(ranked)}, получили {idx}")
    chosen = ranked[idx - 1]

    raw_gb = ask(
        "Data disk GB (тариф 24/7; рост только вверх)",
        str(cfg.data_volume_size_gb),
    ).strip()
    try:
        data_gb = int(raw_gb)
    except ValueError as exc:
        raise ValueError(f"Data disk GB: жду число, получили {raw_gb!r}") from exc
    if data_gb < 20:
        raise ValueError("Data disk GB: минимум 20")

    from gpu_rent.prompts import MenuItem, prompt_menu

    spot_key = prompt_menu(
        "Тариф GPU",
        [
            MenuItem("yes", "preemptible (дешевле, могут усыпить ~24ч)"),
            MenuItem("no", "обычный on-demand"),
        ],
        default="yes" if spot else "no",
        ask=ask,
    )
    use_spot = spot_key == "yes"

    changed = (
        chosen.id != picked.id
        or data_gb != cfg.data_volume_size_gb
        or use_spot != spot
    )
    if confirm and changed and confirm("Запомнить flavor/disk/spot в gpu-rent.vars?"):
        from gpu_rent.paths import vars_path
        from gpu_rent.varsfile import upsert_vars

        upsert_vars(
            vars_path(),
            {
                "DEFAULT_FLAVOR_ID": chosen.id,
                "DATA_VOLUME_SIZE_GB": str(data_gb),
                "DEFAULT_SPOT": "true" if use_spot else "false",
            },
        )

    return ServerPlan(flavor=chosen, data_gb=data_gb, spot=use_spot)


def resolve_and_preview(
    cfg: Config,
    flavors: list[Any],
    *,
    explicit: str | None,
    spot: bool,
    log: Log,
) -> FlavorInfo:
    picked = resolve_flavor(
        flavors,
        cfg.flavor_preference,
        explicit=explicit,
        default_id=cfg.default_flavor_id or None,
        fallback=cfg.flavor_fallback,
    )
    print_up_preview(cfg, flavors, picked=picked, spot=spot, log=log)
    return picked