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 json
import logging
from datetime import datetime, timezone
from typing import Any

from sqlalchemy import and_, or_, select

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).where(*filters).order_by(Server.id)
    )
    rows = [(int(row[0]), row[1] or f"server-{row[0]}") 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,
    }


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