Nexus/server/api/websocket.py

"""Nexus — WebSocket Alert Push (Two-Layer Architecture)

ADR-010: Memory ConnectionManager + Redis Pub/Sub two-layer architecture.
  - Layer 1: In-memory ConnectionManager tracks all local WebSocket connections
  - Layer 2: Redis Pub/Sub relays messages across multiple uvicorn workers
  - Alert/recovery/system events published to Redis channel `nexus:alerts`
  - Each worker subscribes to `nexus:alerts` and pushes to local clients

ADR-011: WebSocket JWT authentication via query parameter `?token=xxx`

R4: Ping/pong heartbeat every 30s + zombie connection cleanup (60s idle → disconnect)
"""

import asyncio
import json
import logging
import time
from typing import Dict, Set, Optional

from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query

from server.config import settings

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

router = APIRouter()

# Redis Pub/Sub channel name
REDIS_CHANNEL = "nexus:alerts"


class ConnectionManager:
    """Layer 1: In-memory WebSocket connection manager.

    Tracks all connected clients for this worker process.
    Handles ping/pong heartbeat and zombie connection cleanup.
    """

    def __init__(self):
        # client_id → {websocket, last_pong_time}
        self._connections: Dict[str, dict] = {}
        self._heartbeat_task: Optional[asyncio.Task] = None

    async def connect(self, client_id: str, websocket: WebSocket):
        """Accept and register a new WebSocket connection"""
        await websocket.accept()
        self._connections[client_id] = {
            "websocket": websocket,
            "last_pong": time.monotonic(),
        }
        logger.info(f"WebSocket connected: {client_id}, total: {len(self._connections)}")

        # Start heartbeat task if first connection
        if len(self._connections) == 1 and self._heartbeat_task is None:
            self._heartbeat_task = asyncio.create_task(self._heartbeat_loop(), name="ws_heartbeat")

    def disconnect(self, client_id: str):
        """Remove a WebSocket connection"""
        if client_id in self._connections:
            del self._connections[client_id]
            logger.info(f"WebSocket disconnected: {client_id}, total: {len(self._connections)}")

        # Stop heartbeat task if no connections
        if not self._connections and self._heartbeat_task:
            self._heartbeat_task.cancel()
            self._heartbeat_task = None

    async def send_to_client(self, client_id: str, message: dict) -> bool:
        """Send JSON message to a specific client. Returns False if send failed."""
        conn = self._connections.get(client_id)
        if not conn:
            return False
        try:
            await conn["websocket"].send_json(message)
            return True
        except Exception:
            self.disconnect(client_id)
            return False

    async def broadcast_local(self, message: dict):
        """Broadcast message to all locally connected clients (Layer 1 only)"""
        stale = []
        for client_id, conn in self._connections.items():
            try:
                await conn["websocket"].send_json(message)
            except Exception:
                stale.append(client_id)
        for client_id in stale:
            self.disconnect(client_id)

    def update_pong(self, client_id: str):
        """Update last pong time for a client (received pong response)"""
        conn = self._connections.get(client_id)
        if conn:
            conn["last_pong"] = time.monotonic()

    @property
    def client_count(self) -> int:
        return len(self._connections)

    async def _heartbeat_loop(self):
        """R4: Send ping every 30s, cleanup connections idle > 60s"""
        while True:
            await asyncio.sleep(30)
            now = time.monotonic()
            stale = []

            for client_id, conn in self._connections.items():
                # Check for zombie connections (no pong for 60s)
                if now - conn["last_pong"] > 60:
                    stale.append(client_id)
                    continue

                # Send ping
                try:
                    await conn["websocket"].send_json({"type": "ping"})
                except Exception:
                    stale.append(client_id)

            # Cleanup stale connections
            for client_id in stale:
                self.disconnect(client_id)
                logger.info(f"WebSocket zombie cleanup: {client_id}")

            if stale:
                logger.info(f"Heartbeat cleanup: removed {len(stale)} stale connections")


# Global connection manager for this worker
manager = ConnectionManager()


# ── Redis Pub/Sub Layer 2 ──

_pubsub_task: Optional[asyncio.Task] = None
_redis_subscriber = None


async def start_redis_subscriber():
    """Start Redis Pub/Sub subscriber (Layer 2) — called during app startup"""
    global _pubsub_task, _redis_subscriber

    try:
        from server.infrastructure.redis.client import get_redis
        redis = get_redis()
        _redis_subscriber = redis.pubsub()
        await _redis_subscriber.subscribe(REDIS_CHANNEL)

        _pubsub_task = asyncio.create_task(_redis_listen_loop(), name="ws_redis_sub")
        logger.info(f"Redis Pub/Sub subscriber started on channel: {REDIS_CHANNEL}")
    except Exception as e:
        logger.error(f"Redis Pub/Sub subscriber failed to start: {e}")
        # Non-fatal — WebSocket still works locally


async def stop_redis_subscriber():
    """Stop Redis Pub/Sub subscriber — called during app shutdown"""
    global _pubsub_task, _redis_subscriber

    if _pubsub_task:
        _pubsub_task.cancel()
        try:
            await _pubsub_task
        except asyncio.CancelledError:
            pass
        _pubsub_task = None

    if _redis_subscriber:
        try:
            await _redis_subscriber.unsubscribe(REDIS_CHANNEL)
            await _redis_subscriber.aclose()
        except Exception:
            logger.debug("Failed to unsubscribe Redis Pub/Sub during shutdown", exc_info=True)
        _redis_subscriber = None

    logger.info("Redis Pub/Sub subscriber stopped")


async def _redis_listen_loop():
    """Listen for messages on Redis Pub/Sub channel and broadcast to local clients"""
    while True:
        try:
            message = await _redis_subscriber.get_message(
                ignore_subscribe_messages=True, timeout=1.0
            )
            if message and message.get("type") == "message":
                data = message.get("data")
                if isinstance(data, str):
                    msg = json.loads(data)
                    await manager.broadcast_local(msg)
        except asyncio.CancelledError:
            break
        except Exception as e:
            logger.error(f"Redis Pub/Sub listen error: {e}")
            await asyncio.sleep(1)  # Back off on error


async def _publish_to_redis(message: dict) -> bool:
    """Publish message to Redis channel (Layer 2) — other workers will receive it"""
    try:
        from server.infrastructure.redis.client import get_redis
        redis = get_redis()
        await redis.publish(REDIS_CHANNEL, json.dumps(message))
        return True
    except Exception as e:
        logger.error(f"Redis Pub/Sub publish failed: {e}")
        return False


async def _dispatch_ws_message(message: dict) -> None:
    """Deliver alert to local WebSocket clients without duplicate delivery.

    When Pub/Sub is active, only publish — this worker's subscriber calls broadcast_local once.
    Fallback to direct broadcast if Redis subscriber is unavailable.
    """
    if _pubsub_task is not None and _redis_subscriber is not None:
        if not await _publish_to_redis(message):
            await manager.broadcast_local(message)
    else:
        await manager.broadcast_local(message)


# ── WebSocket Endpoint ──

@router.websocket("/ws/alerts")
async def alert_ws(
    websocket: WebSocket,
    token: Optional[str] = Query(None, description="JWT access token"),
):
    """WebSocket endpoint for real-time alert push to frontend

    ADR-011: JWT authentication via query parameter `?token=xxx`
    R4: Ping/pong heartbeat + zombie connection cleanup
    """
    # ── JWT Authentication (ADR-011) ──
    if not token:
        await websocket.close(code=4001, reason="Missing JWT token")
        return

    admin = await _verify_ws_token(token)
    if not admin:
        await websocket.close(code=4001, reason="Invalid or expired JWT token")
        return

    # ── Accept and register ──
    client_id = f"{admin.id}:{id(websocket)}"
    await manager.connect(client_id, websocket)

    try:
        while True:
            data = await websocket.receive_text()
            if data == "pong":
                manager.update_pong(client_id)
            elif data == "ping":
                await websocket.send_json({"type": "pong"})
    except WebSocketDisconnect:
        manager.disconnect(client_id)
    except Exception as e:
        logger.warning(f"WebSocket error for {client_id}: {e}")
        manager.disconnect(client_id)


async def _verify_ws_token(token: str):
    """Verify JWT token for WebSocket connection (same logic as get_current_admin)

    Checks: valid signature, not expired, sub exists, admin active, and
    updated_at matches (invalidates token after password change).
    """
    try:
        import jwt as pyjwt
        from server.infrastructure.database.session import AsyncSessionLocal
        from server.infrastructure.database.admin_repo import AdminRepositoryImpl

        payload = pyjwt.decode(token, settings.SECRET_KEY, algorithms=["HS256"],
                                options={"require": ["exp", "sub"]})
        admin_id = payload.get("sub")
        if not admin_id:
            return None

        try:
            admin_id = int(admin_id)
        except (ValueError, TypeError):
            return None

        async with AsyncSessionLocal() as session:
            admin_repo = AdminRepositoryImpl(session)
            admin = await admin_repo.get_by_id(admin_id)
            if not admin or not admin.is_active:
                return None

            # P1-7: Check updated_at — invalidate token after password change
            token_updated = payload.get("updated", 0)
            if token_updated and admin.updated_at:
                admin_ts = int(admin.updated_at.timestamp())
                # Allow 5-second grace for clock skew / race conditions
                if admin_ts > token_updated + 5:
                    return None

            return admin
    except Exception:
        logger.debug("WebSocket JWT verification failed", exc_info=True)
        return None


# ── Alert Aggregation / Dedup ──
# Track last alert time per (server_id, alert_type) to avoid Telegram alert storms.
# WebSocket broadcasts are always sent (real-time), but Telegram pushes are deduplicated.
_ALERT_COOLDOWN_SECONDS = 300  # 5 minutes — same server+metric won't re-push Telegram
_last_alert_time: Dict[str, float] = {}  # key="alert:{server_id}:{type}" → monotonic timestamp


def _should_push_telegram(alert_key: str) -> bool:
    """Check if enough time has passed since last Telegram push for this alert key.

    Returns True if we should send the Telegram notification, False if suppressed (cooldown).
    """
    now = time.monotonic()
    last = _last_alert_time.get(alert_key, 0)
    if now - last < _ALERT_COOLDOWN_SECONDS:
        return False  # Suppressed — too soon
    _last_alert_time[alert_key] = now
    return True


# ── Broadcast Functions (called from other modules) ──

async def _save_alert_log(server_id: int, server_name: str, alert_type: str, value: str, is_recovery: bool):
    """Persist alert / recovery event to MySQL alert_logs table."""
    try:
        from server.infrastructure.database.session import AsyncSessionLocal
        from server.domain.models import AlertLog
        async with AsyncSessionLocal() as session:
            session.add(AlertLog(
                server_id=server_id,
                server_name=server_name or None,
                alert_type=alert_type,
                value=value,
                is_recovery=is_recovery,
            ))
            await session.commit()
    except Exception as e:
        logger.debug("Failed to save alert log: %s", e)


async def broadcast_alert(server_id: int, alert_type: str, alert_value: float, server_name: str = ""):
    """Broadcast alert to all WebSocket clients + Telegram

    Called when Agent heartbeat detects CPU/memory/disk > threshold.
    Publishes to Redis Pub/Sub (multi-worker) + Telegram.

    WebSocket refreshes dashboard stats only (no browser sound/list UI).
    Telegram push is deduplicated (5min cooldown per server+metric).
    """
    msg = {
        "type": "alert",
        "server_id": server_id,
        "server_name": server_name,
        "alert_type": alert_type,
        "alert_value": alert_value,
    }

    # Redis Pub/Sub (multi-worker) — local clients via subscriber only (C9)
    await _dispatch_ws_message(msg)

    # Persist to alert_logs
    await _save_alert_log(server_id, server_name, alert_type, f"{alert_value:.1f}%", is_recovery=False)

    # Telegram push — deduplicated per (server_id, alert_type) with cooldown
    alert_key = f"alert:{server_id}:{alert_type}"
    if _should_push_telegram(alert_key):
        from server.infrastructure.telegram import send_telegram_alert
        await send_telegram_alert(server_name or str(server_id), alert_type, alert_value)
    else:
        logger.debug(f"Alert suppressed (cooldown): {alert_key}")


async def broadcast_recovery(server_id: int, metric: str, value: float, server_name: str = ""):
    """Broadcast recovery notification when alert metric returns to normal

    Recovery alerts are always pushed (clears the alert state).
    Also clears the alert cooldown so next alert will be sent immediately.
    """
    msg = {
        "type": "recovery",
        "server_id": server_id,
        "server_name": server_name,
        "metric": metric,
        "value": value,
    }

    await _dispatch_ws_message(msg)

    # Persist recovery to alert_logs
    await _save_alert_log(server_id, server_name, metric, f"{value:.1f}%", is_recovery=True)

    # Recovery always sends Telegram + clears cooldown for next alert cycle
    alert_key = f"alert:{server_id}:{metric}"
    _last_alert_time.pop(alert_key, None)  # Clear cooldown

    from server.infrastructure.telegram import send_telegram_recovery
    await send_telegram_recovery(server_name or str(server_id), metric, value)


async def broadcast_system_event(event_type: str, message: str):
    """Broadcast system-level events (Python restart, Redis reconnect, etc.)"""
    msg = {"type": "system", "event_type": event_type, "message": message}

    await _dispatch_ws_message(msg)