Nexus/server/application/server_connectivity.py

"""Agent connectivity counts and filters — Redis overlay per DB server (not global key scan)."""

from __future__ import annotations

import asyncio
import json
import logging
from datetime import datetime, timezone
from typing import Any

from sqlalchemy import and_, or_, select

from server.application.server_batch_common import batch_server_display_name
from server.domain.models import Server
from server.infrastructure.database.server_repo import _target_path_unset_filter

logger = logging.getLogger("nexus.server_connectivity")

REDIS_KEY_PREFIX = "heartbeat:"


def heartbeat_indicates_online(heartbeat: dict[str, Any] | None) -> bool:
    """True when Redis heartbeat hash exists and is_online is True (same as stats/list)."""
    return bool(heartbeat) and heartbeat.get("is_online") == "True"


def monitored_server_filter():
    """Servers expected to run Nexus Agent (API key or version on file)."""
    return or_(
        and_(Server.agent_api_key.isnot(None), Server.agent_api_key != ""),
        and_(Server.agent_version.isnot(None), Server.agent_version != ""),
    )


def item_matches_online_filter(item: dict[str, Any], want_online: bool) -> bool:
    """Match UI status chip: online vs offline (excludes unknown)."""
    status = item.get("status") or "unknown"
    if want_online:
        return status == "online"
    return status == "offline"


async def count_monitored_connectivity(
    redis: Any,
    session: Any,
    *,
    exclude_unset_target_path: bool = False,
    batch_size: int = 200,
) -> dict[str, int]:
    """Count online/offline among Agent-monitored servers using per-id Redis heartbeats."""
    scan = await scan_monitored_connectivity(
        redis,
        session,
        exclude_unset_target_path=exclude_unset_target_path,
        batch_size=batch_size,
    )
    return {
        "monitored": scan["monitored"],
        "online": scan["online"],
        "offline": scan["offline"],
    }


async def scan_monitored_connectivity(
    redis: Any,
    session: Any,
    *,
    exclude_unset_target_path: bool = False,
    batch_size: int = 200,
) -> dict[str, Any]:
    """Count connectivity and list offline Agent-monitored servers."""
    filters = [monitored_server_filter()]
    if exclude_unset_target_path:
        filters.append(_target_path_unset_filter(False))

    result = await session.execute(
        select(Server.id, Server.name, Server.domain, Server.description)
        .where(*filters)
        .order_by(Server.id)
    )
    rows = [
        (
            int(row[0]),
            batch_server_display_name(
                Server(id=int(row[0]), name=row[1], domain=row[2], description=row[3])
            ),
        )
        for row in result.all()
    ]
    monitored = len(rows)
    if not monitored:
        return {"monitored": 0, "online": 0, "offline": 0, "offline_servers": []}

    online = 0
    offline = 0
    offline_servers: list[dict[str, Any]] = []

    for i in range(0, monitored, batch_size):
        chunk = rows[i : i + batch_size]
        pipe = redis.pipeline()
        for server_id, _ in chunk:
            pipe.hgetall(f"{REDIS_KEY_PREFIX}{server_id}")
        heartbeats = await pipe.execute()

        for (server_id, server_name), hb in zip(chunk, heartbeats):
            if heartbeat_indicates_online(hb):
                online += 1
            else:
                offline += 1
                offline_servers.append({
                    "id": server_id,
                    "name": server_name,
                    "last_heartbeat": (hb or {}).get("last_heartbeat", ""),
                })

    return {
        "monitored": monitored,
        "online": online,
        "offline": offline,
        "offline_servers": offline_servers,
    }


async def health_check_server_ids(
    server_ids: list[int],
    *,
    concurrency: int = 10,
    timeout: int = 15,
) -> dict[str, int]:
    """SSH health-check given servers and refresh Redis heartbeats (batch health-check logic)."""
    ids = [int(sid) for sid in server_ids if sid]
    if not ids:
        return {"checked": 0, "online": 0, "offline": 0, "errors": 0}

    from server.infrastructure.database.server_repo import ServerRepositoryImpl
    from server.infrastructure.database.session import AsyncSessionLocal
    from server.infrastructure.redis.client import get_redis
    from server.infrastructure.ssh.remote_probe import ssh_health_probe

    redis = get_redis()
    sem = asyncio.Semaphore(max(1, concurrency))
    stats = {"checked": 0, "online": 0, "offline": 0, "errors": 0}

    async with AsyncSessionLocal() as session:
        repo = ServerRepositoryImpl(session)
        servers = await repo.get_by_ids(ids)
    server_map = {s.id: s for s in servers}

    async def _one(sid: int) -> None:
        async with sem:
            server = server_map.get(sid)
            if not server:
                stats["errors"] += 1
                return
            try:
                probe = await ssh_health_probe(server, timeout=timeout)
                await write_ssh_health_heartbeat(redis, server, probe)
                if probe.get("status") == "online":
                    stats["online"] += 1
                else:
                    stats["offline"] += 1
            except Exception as exc:
                stats["errors"] += 1
                logger.warning("offline preflight health check failed server=%s: %s", sid, exc)
            finally:
                stats["checked"] += 1

    await asyncio.gather(*[_one(sid) for sid in ids])
    return stats


def normalize_ssh_probe_system_info(probe: dict[str, Any]) -> dict[str, Any]:
    """Flatten ssh_health_probe payload into Agent-compatible system_info."""
    raw = probe.get("system_info") or {}
    if isinstance(raw, dict) and isinstance(raw.get("system_info"), dict):
        system_info = dict(raw["system_info"])
    elif isinstance(raw, dict):
        system_info = dict(raw)
    else:
        system_info = {}
    system_info["probe_channel"] = str(probe.get("channel") or "ssh")
    if probe.get("status") != "online" and probe.get("error"):
        system_info["probe_error"] = str(probe["error"])[:300]
    return system_info


async def write_ssh_health_heartbeat(
    redis: Any,
    server: Server,
    probe: dict[str, Any],
) -> str:
    """Write SSH health-check result to Redis heartbeat (same shape as Agent). Returns last_heartbeat ISO."""
    from server.background.heartbeat_flush import FLUSH_INTERVAL

    is_online = probe.get("status") == "online"
    system_info = normalize_ssh_probe_system_info(probe)
    now = datetime.now(timezone.utc).isoformat()
    agent_version = (server.agent_version or "").strip()

    key = f"{REDIS_KEY_PREFIX}{server.id}"
    await redis.hset(key, mapping={
        "is_online": str(is_online),
        "system_info": json.dumps(system_info, ensure_ascii=False),
        "last_heartbeat": now,
        "agent_version": agent_version,
        "time_drift_seconds": "0",
        "drift_level": "ok",
    })
    await redis.expire(key, int(FLUSH_INTERVAL * 1.5))
    return now