Nexus/server/application/services/btpanel_bootstrap_schedule.py

"""Fire-and-forget and background tick for Baota SSH API bootstrap."""

from __future__ import annotations

import asyncio
import logging
from datetime import datetime, timezone

from server.config import settings
from server.infrastructure.btpanel.bootstrap_state import (
    is_eligible_for_background_bootstrap,
    mark_bootstrap_pending,
    _parse_iso,
)
from server.infrastructure.btpanel.credentials import bt_panel_configured
from server.infrastructure.database.server_repo import ServerRepositoryImpl
from server.infrastructure.database.session import AsyncSessionLocal

logger = logging.getLogger("nexus.btpanel.bootstrap_schedule")

_BG_TASKS: set[asyncio.Task] = set()


def is_auto_bootstrap_enabled() -> bool:
    val = (getattr(settings, "BT_PANEL_AUTO_BOOTSTRAP_ENABLED", None) or "true").strip().lower()
    return val not in ("false", "0", "no", "off")


def bootstrap_batch_size() -> int:
    raw = getattr(settings, "BT_PANEL_BOOTSTRAP_BATCH", 20)
    try:
        n = int(raw)
    except (TypeError, ValueError):
        n = 20
    return max(1, min(n, 100))


async def mark_servers_bootstrap_pending(server_ids: list[int]) -> None:
    ids = [int(x) for x in server_ids if x]
    if not ids:
        return
    async with AsyncSessionLocal() as session:
        repo = ServerRepositoryImpl(session)
        for sid in ids:
            server = await repo.get_by_id(sid)
            if not server or bt_panel_configured(server):
                continue
            server.extra_attrs = mark_bootstrap_pending(
                server.extra_attrs if isinstance(server.extra_attrs, dict) else {},
            )
            await repo.update(server)


async def schedule_immediate_bootstrap(server_ids: list[int], *, source: str = "immediate") -> None:
    """Mark pending and try once without waiting for the 5-minute loop."""
    if not is_auto_bootstrap_enabled():
        return
    ids = [int(x) for x in server_ids if x]
    if not ids:
        return
    await mark_servers_bootstrap_pending(ids)
    for sid in ids:
        task = asyncio.create_task(_bootstrap_one_background(sid, source=source))
        _BG_TASKS.add(task)
        task.add_done_callback(_BG_TASKS.discard)


async def _bootstrap_one_background(server_id: int, *, source: str) -> None:
    from server.application.services.btpanel_service import BtPanelService

    try:
        async with AsyncSessionLocal() as session:
            svc = BtPanelService(session)
            await svc.bootstrap_server(
                server_id,
                source=source,
                admin_username="system",
                ip_address=None,
            )
    except Exception as exc:
        logger.warning("bt panel bootstrap task failed server=%s source=%s: %s", server_id, source, exc)


async def run_bootstrap_tick() -> int:
    """Pick due servers and bootstrap up to batch limit; returns attempt count."""
    if not is_auto_bootstrap_enabled():
        return 0

    from server.application.services.btpanel_service import BtPanelService

    now = datetime.now(timezone.utc)
    limit = bootstrap_batch_size()
    due_ids: list[int] = []

    async with AsyncSessionLocal() as session:
        repo = ServerRepositoryImpl(session)
        servers = await repo.get_all()
        candidates: list[tuple[datetime, int]] = []
        for server in servers:
            if is_eligible_for_background_bootstrap(server, now=now):
                from server.infrastructure.btpanel.credentials import get_bt_panel_block

                block = get_bt_panel_block(server)
                next_raw = block.get("bootstrap_next_attempt_at") if block else None
                next_at = _parse_iso(next_raw) or now
                candidates.append((next_at, server.id))
        candidates.sort(key=lambda x: x[0])
        due_ids = [sid for _, sid in candidates[:limit]]

    if not due_ids:
        return 0

    succeeded = 0
    for sid in due_ids:
        try:
            async with AsyncSessionLocal() as session:
                svc = BtPanelService(session)
                result = await svc.bootstrap_server(
                    sid,
                    source="background",
                    admin_username="system",
                    ip_address=None,
                )
            if result.get("ok"):
                succeeded += 1
        except Exception as exc:
            logger.warning("bt panel background bootstrap failed server=%s: %s", sid, exc)
    return succeeded