gpu-rent/src/gpu_rent/balance.py

"""Selectel account balance watch: notify every N ₽ spent since up."""

from __future__ import annotations

import json
from collections.abc import Callable
from dataclasses import asdict, dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import httpx

from gpu_rent.errors import CloudError
from gpu_rent.paths import runtime_dir

Log = Callable[[str], None]

BALANCES_URL = "https://api.selectel.ru/v3/balances"
DEFAULT_STEP_RUB = 200.0


@dataclass
class BalanceWatchState:
    baseline_rub: float
    last_notified_step: int = 0
    low_notified: bool = False
    last_balance_rub: float | None = None
    updated_at: str | None = None

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> BalanceWatchState:
        return cls(
            baseline_rub=float(data.get("baseline_rub") or 0),
            last_notified_step=int(data.get("last_notified_step") or 0),
            low_notified=bool(data.get("low_notified")),
            last_balance_rub=(
                float(data["last_balance_rub"])
                if data.get("last_balance_rub") is not None
                else None
            ),
            updated_at=data.get("updated_at"),
        )


def balance_watch_path() -> Path:
    return runtime_dir() / "balance-watch.json"


def utc_now_iso() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat()


def load_balance_state() -> BalanceWatchState | None:
    path = balance_watch_path()
    if not path.is_file():
        return None
    try:
        raw = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return None
    if not isinstance(raw, dict) or "baseline_rub" not in raw:
        return None
    return BalanceWatchState.from_dict(raw)


def save_balance_state(state: BalanceWatchState) -> None:
    runtime_dir().mkdir(parents=True, exist_ok=True)
    state.updated_at = utc_now_iso()
    balance_watch_path().write_text(
        json.dumps(state.to_dict(), indent=2) + "\n", encoding="utf-8"
    )


def clear_balance_state() -> None:
    path = balance_watch_path()
    if path.is_file():
        path.unlink()


def balance_rub_from_payload(payload: dict[str, Any]) -> float:
    """Selectel returns money as integer kopecks → rubles."""
    data = payload.get("data") if isinstance(payload.get("data"), dict) else payload
    billings = data.get("billings") if isinstance(data, dict) else None
    if not isinstance(billings, list) or not billings:
        raise CloudError("Selectel balances: нет data.billings")
    total_kop = 0
    for row in billings:
        if not isinstance(row, dict):
            continue
        if row.get("balances_values_sum") is not None:
            total_kop += int(row["balances_values_sum"])
        elif row.get("final_sum") is not None:
            total_kop += int(row["final_sum"])
    return total_kop / 100.0


def fetch_balance_rub(api_token: str, *, timeout: float = 20.0) -> float:
    token = (api_token or "").strip()
    if not token:
        raise CloudError("SELECTEL_API_TOKEN пуст — баланс недоступен")
    headers = {"X-Token": token, "Accept": "application/json"}
    try:
        with httpx.Client(timeout=timeout, follow_redirects=True) as client:
            resp = client.get(BALANCES_URL, headers=headers)
    except httpx.HTTPError as exc:
        raise CloudError(f"Selectel balances сеть: {exc}") from exc
    if resp.status_code in {401, 403}:
        raise CloudError(
            f"Selectel balances HTTP {resp.status_code} — нужен статический "
            "API-ключ аккаунта (X-Token), не OpenStack пароль"
        )
    if resp.status_code != 200:
        raise CloudError(f"Selectel balances HTTP {resp.status_code}: {resp.text[:200]}")
    try:
        payload = resp.json()
    except json.JSONDecodeError as exc:
        raise CloudError("Selectel balances: не JSON") from exc
    if not isinstance(payload, dict):
        raise CloudError("Selectel balances: неожиданный JSON")
    return balance_rub_from_payload(payload)


def spend_steps(baseline_rub: float, current_rub: float, step_rub: float) -> int:
    """How many full step intervals have been spent since baseline."""
    if step_rub <= 0:
        return 0
    spend = baseline_rub - current_rub
    if spend < 0:
        return 0
    return int(spend // step_rub)


def plan_balance_notices(
    state: BalanceWatchState,
    current_rub: float,
    *,
    step_rub: float = DEFAULT_STEP_RUB,
    low_rub: float = 0.0,
) -> tuple[BalanceWatchState, list[str]]:
    """
    Update state for top-ups / spend steps / low watermark.
    Returns (new_state, human messages to show).
    """
    messages: list[str] = []
    # Top-up: reset baseline so steps restart from new balance.
    if current_rub > state.baseline_rub + 0.01:
        state = BalanceWatchState(
            baseline_rub=current_rub,
            last_notified_step=0,
            low_notified=False,
            last_balance_rub=current_rub,
        )
        messages.append(f"Баланс пополнен: {current_rub:.0f} ₽ (новый baseline)")
        return state, messages

    crossed = spend_steps(state.baseline_rub, current_rub, step_rub)
    while state.last_notified_step < crossed:
        state.last_notified_step += 1
        spent = state.last_notified_step * step_rub
        messages.append(
            f"Списано ~{spent:.0f} ₽ с момента up "
            f"(осталось {current_rub:.0f} ₽, шаг {step_rub:.0f} ₽)"
        )

    if low_rub > 0 and current_rub < low_rub and not state.low_notified:
        state.low_notified = True
        messages.append(f"Баланс низкий: {current_rub:.0f} ₽ (< {low_rub:.0f} ₽)")

    state.last_balance_rub = current_rub
    return state, messages