"""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, Optional

from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query

from server.config import settings
from server.utils.display_time import utc_now

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

router = APIRouter()

# Redis Pub/Sub channel names (alerts vs push progress — separate to avoid cross-traffic)
REDIS_CHANNEL = "nexus:alerts"
REDIS_SYNC_CHANNEL = "nexus:sync"
REDIS_WATCH_CHANNEL = "nexus:watch"


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.
    """

    MAX_CONNECTIONS = 500

    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"""
        if len(self._connections) >= self.MAX_CONNECTIONS:
            logger.warning(f"WebSocket connection limit reached ({self.MAX_CONNECTIONS}), rejecting {client_id}")
            await websocket.close(code=4003, reason="Too many connections")
            return
        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 and close the underlying TCP connection"""
        conn = self._connections.pop(client_id, None)
        if conn:
            # Close the WebSocket to release the TCP connection
            try:
                asyncio.get_running_loop().create_task(
                    conn["websocket"].close(code=1000, reason="cleanup")
                )
            except RuntimeError:
                pass  # No running event loop (e.g. during shutdown)
            except Exception:  # noqa: S110 — client disconnected, best-effort cleanup
                pass
            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:
        """Local client count (this worker only)"""
        return len(self._connections)

    async def global_client_count(self) -> int:
        """Global client count across all workers (via Redis).

        In multi-worker deployments, each worker only knows its local connections.
        This method queries Redis for the aggregate count.
        """
        try:
            from server.infrastructure.redis.client import get_redis
            redis = get_redis()
            # Each worker periodically writes its count to Redis with TTL
            await redis.set(f"ws:count:{id(self)}", len(self._connections), ex=90)
            # Sum all worker counts — use SCAN instead of KEYS to avoid blocking
            total = 0
            async for key in redis.scan_iter(match="ws:count:*"):
                val = await redis.get(key)
                if val:
                    total += int(val)
            return total
        except Exception:
            return len(self._connections)

    async def _heartbeat_loop(self):
        """R4: Send ping every 30s, cleanup connections idle > 60s"""
        try:
            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")
        except asyncio.CancelledError:
            logger.debug("WebSocket heartbeat task cancelled")
            raise

    def cancel_heartbeat(self):
        """Cancel the heartbeat task — called during app shutdown (H11)"""
        if self._heartbeat_task:
            self._heartbeat_task.cancel()
            self._heartbeat_task = None


# Global connection managers for this worker
manager = ConnectionManager()
sync_manager = ConnectionManager()
watch_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, REDIS_SYNC_CHANNEL, REDIS_WATCH_CHANNEL)

        _pubsub_task = asyncio.create_task(_redis_listen_loop(), name="ws_redis_sub")
        logger.info(
            "Redis Pub/Sub subscriber started on channels: %s, %s, %s",
            REDIS_CHANNEL,
            REDIS_SYNC_CHANNEL,
            REDIS_WATCH_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, REDIS_SYNC_CHANNEL, REDIS_WATCH_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")
                channel = message.get("channel")
                if isinstance(channel, bytes):
                    channel = channel.decode()
                if isinstance(data, str):
                    msg = json.loads(data)
                    if channel == REDIS_SYNC_CHANNEL:
                        await sync_manager.broadcast_local(msg)
                    elif channel == REDIS_WATCH_CHANNEL:
                        await _broadcast_watch_local(msg)
                    else:
                        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, channel: str = REDIS_CHANNEL) -> 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(channel, json.dumps(message))
        return True
    except Exception as e:
        logger.error(f"Redis Pub/Sub publish failed on {channel}: {e}")
        return False


async def _dispatch_sync_message(message: dict) -> None:
    """Deliver sync progress to /ws/sync clients only (separate from alert channel)."""
    if _pubsub_task is not None and _redis_subscriber is not None:
        if not await _publish_to_redis(message, REDIS_SYNC_CHANNEL):
            await sync_manager.broadcast_local(message)
    else:
        await sync_manager.broadcast_local(message)


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)


async def _broadcast_watch_local(message: dict) -> None:
    """Push watch_metrics to local /ws/watch clients filtered by admin_id."""
    admin_ids = message.get("admin_ids") or []
    if not admin_ids:
        return
    admin_set = {int(a) for a in admin_ids}
    payload = {
        "type": "watch_metrics",
        "server_id": message.get("server_id"),
        "metrics": message.get("metrics") or {},
    }
    stale: list[str] = []
    for client_id, conn in watch_manager._connections.items():
        parts = client_id.split(":", 2)
        if len(parts) < 2:
            continue
        try:
            admin_id = int(parts[1])
        except ValueError:
            continue
        if admin_id not in admin_set:
            continue
        try:
            await conn["websocket"].send_json(payload)
        except Exception:
            stale.append(client_id)
    for client_id in stale:
        watch_manager.disconnect(client_id)


async def publish_watch_metrics(message: dict) -> None:
    """Publish watch probe update to Redis; local clients via subscriber."""
    if _pubsub_task is not None and _redis_subscriber is not None:
        if not await _publish_to_redis(message, REDIS_WATCH_CHANNEL):
            await _broadcast_watch_local(message)
    else:
        await _broadcast_watch_local(message)


@router.websocket("/ws/watch")
async def watch_metrics_ws(
    websocket: WebSocket,
    token: Optional[str] = Query(None, description="JWT access token"),
):
    """WebSocket for real-time watch slot metrics (5s probe updates)."""
    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=4401, reason="Invalid or expired JWT token")
        return

    client_id = f"watch:{admin.id}:{id(websocket)}"
    await watch_manager.connect(client_id, websocket)

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


# ── 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=4401, 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)


@router.websocket("/ws/sync")
async def sync_progress_ws(
    websocket: WebSocket,
    token: Optional[str] = Query(None, description="JWT access token"),
):
    """WebSocket endpoint dedicated to file-push progress (sync_progress only).

    Separated from /ws/alerts so concurrent pushes and alert traffic do not share a channel.
    """
    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=4401, reason="Invalid or expired JWT token")
        return

    client_id = f"sync:{admin.id}:{id(websocket)}"
    await sync_manager.connect(client_id, websocket)

    try:
        while True:
            data = await websocket.receive_text()
            if data == "pong":
                sync_manager.update_pong(client_id)
            elif data == "ping":
                await websocket.send_json({"type": "pong"})
    except WebSocketDisconnect:
        sync_manager.disconnect(client_id)
    except Exception as e:
        logger.warning(f"Sync WebSocket error for {client_id}: {e}")
        sync_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

            # Primary: token_version must match (immediate invalidation on password change)
            token_tv = payload.get("tv") or 0
            admin_tv = admin.token_version or 0
            if token_tv != admin_tv:
                return None

            # Secondary: updated_at timestamp (defense in depth)
            token_updated = payload.get("updated", 0)
            if token_updated and admin.updated_at:
                admin_ts = int(admin.updated_at.timestamp())
                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).
    Also prunes stale entries older than 1 hour to prevent unbounded memory growth.
    """
    now = time.monotonic()

    # Prune entries older than 1 hour (prevent unbounded growth)
    stale_keys = [k for k, v in _last_alert_time.items() if now - v > 3600]
    for k in stale_keys:
        del _last_alert_time[k]

    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,
        "at": utc_now().isoformat(),
    }

    # 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
        from server.utils.server_alert_ips import resolve_alert_ips_for_server

        public_ip, private_ip = await resolve_alert_ips_for_server(server_id)
        await send_telegram_alert(
            server_name or str(server_id),
            alert_type,
            alert_value,
            server_id,
            public_ip=public_ip,
            private_ip=private_ip,
        )
    else:
        logger.debug(f"Alert suppressed (cooldown): {alert_key}")


async def broadcast_offline_alert(
    server_id: int,
    server_name: str = "",
    last_heartbeat: str = "",
):
    """Broadcast server offline alert when Agent heartbeat is lost (edge-triggered).

    WebSocket refreshes dashboard stats; Telegram respects notify_alert_offline toggle.
    """
    msg = {
        "type": "alert",
        "server_id": server_id,
        "server_name": server_name,
        "alert_type": "offline",
        "alert_value": 0,
        "last_heartbeat": last_heartbeat or None,
        "at": utc_now().isoformat(),
    }

    await _dispatch_ws_message(msg)

    value_str = last_heartbeat if last_heartbeat else "离线"
    await _save_alert_log(server_id, server_name, "offline", value_str, is_recovery=False)

    alert_key = f"alert:{server_id}:offline"
    if _should_push_telegram(alert_key):
        from server.infrastructure.telegram import send_telegram_offline_alert

        from server.utils.server_alert_ips import resolve_alert_ips_for_server

        public_ip, private_ip = await resolve_alert_ips_for_server(server_id)
        await send_telegram_offline_alert(
            server_name or str(server_id),
            server_id,
            last_heartbeat=last_heartbeat,
            public_ip=public_ip,
            private_ip=private_ip,
        )
    else:
        logger.debug("Offline alert suppressed (cooldown): %s", 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,
        "at": utc_now().isoformat(),
    }

    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
    from server.utils.server_alert_ips import resolve_alert_ips_for_server

    public_ip, private_ip = await resolve_alert_ips_for_server(server_id)
    await send_telegram_recovery(
        server_name or str(server_id),
        metric,
        value,
        server_id,
        public_ip=public_ip,
        private_ip=private_ip,
    )


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)


async def broadcast_script_progress(payload: dict):
    """Broadcast script execution batch progress to /ws/alerts subscribers."""
    msg = {"type": "script_progress", **payload}
    await _dispatch_ws_message(msg)


async def broadcast_script_complete(payload: dict):
    """Broadcast script execution terminal state to /ws/alerts subscribers."""
    msg = {"type": "script_complete", **payload}
    await _dispatch_ws_message(msg)


async def broadcast_sync_progress(
    batch_id: str, server_id: int, server_name: str,
    status: str, completed: int, failed: int, total: int,
    error_message: str = None, duration_seconds: int = None,
    retry_job_id: int = None, cancelled: int = 0,
):
    """Broadcast per-server push progress to /ws/sync subscribers.

    Called from sync_engine_v2._sync_one() after each server state change.
    Frontend Push page connects to /ws/sync and filters by batch_id.
    """
    msg = {
        "type": "sync_progress",
        "batch_id": batch_id,
        "server_id": server_id,
        "server_name": server_name,
        "status": status,
        "completed": completed,
        "failed": failed,
        "total": total,
        "error_message": error_message,
        "duration_seconds": duration_seconds,
        "retry_job_id": retry_job_id,
        "cancelled": cancelled,
    }
    await _dispatch_sync_message(msg)