"""WebSocket Connection Manager for real-time event streaming.

This module implements a high-performance WebSocket manager that bridges
Redis Pub/Sub events to connected WebSocket clients.

Issue #1116: Add WebSocket Connection Manager for Real-Time Events

Performance Considerations:
- O(1) connection lookup via nested dict structure
- Single Redis subscription per zone (not per connection)
- Efficient pattern matching reuses existing fnmatch logic
- Lazy zone subscription (only subscribe when first client connects)
- Background task cleanup on disconnect
"""

import asyncio
import logging
import time
from contextlib import suppress
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any

from fastapi import WebSocket, WebSocketDisconnect

if TYPE_CHECKING:
    from nexus.services.event_bus.protocol import EventBusProtocol
    from nexus.services.event_bus.subscriptions import ReactiveSubscriptionManager
    from nexus.services.event_bus.types import FileEvent

logger = logging.getLogger(__name__)

# Constants
PING_INTERVAL = 25.0  # seconds - slightly less than typical 30s proxy timeout
PING_TIMEOUT = 10.0  # seconds to wait for pong response
MAX_MESSAGE_SIZE = 1024 * 1024  # 1MB max message size


@dataclass
class ConnectionInfo:
    """Metadata for a WebSocket connection."""

    websocket: WebSocket
    zone_id: str
    subscription_id: str | None  # None for zone-wide subscriptions
    user_id: str | None
    connected_at: float = field(default_factory=time.time)
    last_pong: float = field(default_factory=time.time)
    event_types: list[str] = field(default_factory=list)  # Event types to filter
    messages_sent: int = 0
    messages_received: int = 0


@dataclass
class WebSocketStats:
    """Statistics for WebSocket connections."""

    total_connections: int = 0
    total_disconnections: int = 0
    total_messages_sent: int = 0
    total_messages_received: int = 0
    current_connections: int = 0
    connections_by_zone: dict[str, int] = field(default_factory=dict)


class WebSocketManager:
    """Manages WebSocket connections and event broadcasting.

    Thread-safe manager for WebSocket connections with Redis Pub/Sub integration.
    Designed for high performance with efficient data structures and lazy subscriptions.

    Architecture:
        zone_id -> connection_id -> ConnectionInfo
        One Redis subscription task per zone (shared by all connections)

    Example:
        manager = WebSocketManager(event_bus)
        await manager.start()

        # In WebSocket endpoint
        await manager.connect(websocket, zone_id, subscription_id, user_id)
        try:
            await manager.handle_client(websocket, connection_id)
        finally:
            await manager.disconnect(connection_id)
    """

    def __init__(
        self,
        event_bus: "EventBusProtocol | None" = None,
        reactive_manager: "ReactiveSubscriptionManager | None" = None,
    ) -> None:
        """Initialize the WebSocket manager.

        Args:
            event_bus: EventBus for Redis Pub/Sub integration (optional)
            reactive_manager: ReactiveSubscriptionManager for O(1) event
                matching (optional, falls back to legacy pattern scan)
        """
        self._event_bus = event_bus
        self._reactive_manager = reactive_manager
        self._started = False

        # Connection tracking: zone_id -> connection_id -> ConnectionInfo
        self._connections: dict[str, dict[str, ConnectionInfo]] = {}

        # Connection ID to zone mapping for O(1) lookup
        self._connection_to_zone: dict[str, str] = {}

        # Redis subscription tasks per zone
        self._subscription_tasks: dict[str, asyncio.Task[None]] = {}

        # Heartbeat task
        self._heartbeat_task: asyncio.Task[None] | None = None

        # Statistics
        self._stats = WebSocketStats()

        # Lock for connection modifications
        self._lock = asyncio.Lock()

    async def start(self) -> None:
        """Start the WebSocket manager.

        Starts the heartbeat task for connection health monitoring.
        """
        if self._started:
            return

        self._started = True
        self._heartbeat_task = asyncio.create_task(self._heartbeat_loop(), name="ws-heartbeat")
        logger.info("WebSocket manager started")

    async def stop(self) -> None:
        """Stop the WebSocket manager and close all connections."""
        if not self._started:
            return

        self._started = False

        # Cancel heartbeat task
        if self._heartbeat_task:
            self._heartbeat_task.cancel()
            with suppress(asyncio.CancelledError):
                await self._heartbeat_task
            self._heartbeat_task = None

        # Cancel all subscription tasks
        for task in self._subscription_tasks.values():
            task.cancel()
        for task in self._subscription_tasks.values():
            with suppress(asyncio.CancelledError):
                await task
        self._subscription_tasks.clear()

        # Close all connections gracefully
        async with self._lock:
            for zone_connections in self._connections.values():
                for conn_info in zone_connections.values():
                    with suppress(Exception):  # Best-effort close during shutdown
                        await conn_info.websocket.close(code=1001, reason="Server shutting down")
            self._connections.clear()
            self._connection_to_zone.clear()

        logger.info("WebSocket manager stopped")

    async def connect(
        self,
        websocket: WebSocket,
        zone_id: str,
        connection_id: str,
        user_id: str | None = None,
        subscription_id: str | None = None,
        event_types: list[str] | None = None,
    ) -> ConnectionInfo:
        """Register a new WebSocket connection.

        Args:
            websocket: The WebSocket connection
            zone_id: Zone ID for event filtering
            connection_id: Unique connection identifier
            user_id: Optional user ID for logging
            subscription_id: Optional subscription ID for filtering
            event_types: Event types to receive

        Returns:
            ConnectionInfo for the new connection
        """
        await websocket.accept()

        conn_info = ConnectionInfo(
            websocket=websocket,
            zone_id=zone_id,
            subscription_id=subscription_id,
            user_id=user_id,
            event_types=event_types or [],
        )

        async with self._lock:
            # Add to connection tracking
            if zone_id not in self._connections:
                self._connections[zone_id] = {}
            self._connections[zone_id][connection_id] = conn_info
            self._connection_to_zone[connection_id] = zone_id

            # Update stats
            self._stats.total_connections += 1
            self._stats.current_connections += 1
            self._stats.connections_by_zone[zone_id] = (
                self._stats.connections_by_zone.get(zone_id, 0) + 1
            )

            # Start Redis subscription for this zone if needed
            if self._event_bus and zone_id not in self._subscription_tasks and self._started:
                task = asyncio.create_task(
                    self._redis_subscription_loop(zone_id),
                    name=f"ws-redis-{zone_id}",
                )
                self._subscription_tasks[zone_id] = task

        logger.info(
            f"WebSocket connected: connection_id={connection_id}, zone={zone_id}, user={user_id}"
        )

        return conn_info

    async def disconnect(self, connection_id: str) -> None:
        """Remove a WebSocket connection.

        Args:
            connection_id: The connection ID to remove
        """
        async with self._lock:
            zone_id = self._connection_to_zone.pop(connection_id, None)
            if not zone_id:
                return

            if zone_id in self._connections:
                conn_info = self._connections[zone_id].pop(connection_id, None)
                if conn_info:
                    self._stats.total_disconnections += 1
                    self._stats.current_connections -= 1
                    if zone_id in self._stats.connections_by_zone:
                        self._stats.connections_by_zone[zone_id] -= 1
                        if self._stats.connections_by_zone[zone_id] <= 0:
                            del self._stats.connections_by_zone[zone_id]

                # Clean up empty zone
                if not self._connections[zone_id]:
                    del self._connections[zone_id]

                    # Cancel Redis subscription if no more connections for zone
                    if zone_id in self._subscription_tasks:
                        self._subscription_tasks[zone_id].cancel()
                        del self._subscription_tasks[zone_id]

        # Clean up reactive subscriptions for this connection
        if self._reactive_manager:
            await self._reactive_manager.unregister_connection(connection_id)

        logger.info(f"WebSocket disconnected: connection_id={connection_id}")

    async def handle_client(self, websocket: WebSocket, connection_id: str) -> None:
        """Handle messages from a WebSocket client.

        This is the main loop for handling client messages. It processes
        ping/pong for heartbeat and any other client messages.

        Args:
            websocket: The WebSocket connection
            connection_id: The connection ID

        Raises:
            WebSocketDisconnect: When the client disconnects
        """
        try:
            while True:
                data = await websocket.receive_json()
                await self._handle_client_message(connection_id, data)
        except WebSocketDisconnect:
            raise
        except Exception as e:
            logger.warning(f"Error handling client message: {e}")
            raise WebSocketDisconnect(code=1011) from e

    async def _handle_client_message(self, connection_id: str, data: dict[str, Any]) -> None:
        """Process a message from a client.

        Args:
            connection_id: The connection ID
            data: The parsed JSON message
        """
        zone_id = self._connection_to_zone.get(connection_id)
        if not zone_id:
            return

        conn_info = self._connections.get(zone_id, {}).get(connection_id)
        if not conn_info:
            return

        conn_info.messages_received += 1
        self._stats.total_messages_received += 1

        msg_type = data.get("type", "")

        if msg_type == "pong":
            conn_info.last_pong = time.time()
            logger.debug(f"Received pong from {connection_id}")
        elif msg_type == "ping":
            # Client-initiated ping, respond with pong
            await self._send_to_connection(conn_info, {"type": "pong"})
        elif msg_type == "subscribe":
            # Update subscription filters
            if "event_types" in data:
                conn_info.event_types = data["event_types"]
            # Echo back patterns and event_types for client confirmation
            response: dict[str, Any] = {"type": "subscribed"}
            if "patterns" in data:
                response["patterns"] = data["patterns"]
            if "event_types" in data:
                response["event_types"] = data["event_types"]
            await self._send_to_connection(conn_info, response)
        else:
            logger.debug(f"Unknown message type from {connection_id}: {msg_type}")

    async def broadcast_to_zone(self, zone_id: str, event: "FileEvent") -> int:
        """Broadcast an event to all connections for a zone.

        Uses ReactiveSubscriptionManager to send batch_update messages (#1170)
        grouping all affected subscriptions per connection into a single
        atomic message for consistent client state.

        Args:
            zone_id: The zone ID
            event: The file event to broadcast

        Returns:
            Number of connections that received the event
        """
        connections = self._connections.get(zone_id, {})
        if not connections:
            return 0

        if not self._reactive_manager:
            return 0

        return await self._broadcast_batch(connections, event)

    async def _broadcast_batch(
        self,
        connections: dict[str, ConnectionInfo],
        event: "FileEvent",
    ) -> int:
        """Send batch_update messages with subscriptions grouped per connection.

        Each connection receives a single batch_update message containing all
        affected subscription IDs, ensuring consistent client state (#1170).

        Args:
            connections: Zone connections dict (connection_id -> ConnectionInfo)
            event: The file event to broadcast

        Returns:
            Number of connections that received the batch
        """
        if self._reactive_manager is None:
            return 0

        try:
            affected_subs = self._reactive_manager.find_affected_subscriptions(event)
        except Exception as e:
            logger.error(f"Reactive subscription lookup failed: {e}")
            return 0

        # Build per-connection batch messages
        event_data = event.to_dict()
        targets: list[tuple[str, ConnectionInfo, dict[str, Any]]] = []

        for conn_id, subs in affected_subs.items():
            conn_info = connections.get(conn_id)
            if not conn_info or not subs:
                continue

            message: dict[str, Any] = {
                "type": "batch_update",
                "event": event_data,
                "commit_id": None,
                "timestamp": event.timestamp,
                "updates": [
                    {
                        "subscription_id": sub.subscription_id,
                        **({"query_id": sub.query_id} if sub.query_id else {}),
                    }
                    for sub in subs
                ],
            }
            targets.append((conn_id, conn_info, message))

        return await self._send_and_cleanup(targets)

    async def _send_and_cleanup(
        self,
        targets: list[tuple[str, ConnectionInfo, dict[str, Any]]],
    ) -> int:
        """Send messages to targets and clean up failed connections.

        Args:
            targets: List of (connection_id, conn_info, message) tuples

        Returns:
            Count of successful sends
        """
        sent_count = 0
        failed_connections: list[str] = []

        for conn_id, conn_info, message in targets:
            try:
                await self._send_to_connection(conn_info, message)
                sent_count += 1
            except Exception as e:
                logger.warning(f"Failed to send to {conn_id}: {e}")
                failed_connections.append(conn_id)

        for conn_id in failed_connections:
            await self.disconnect(conn_id)

        return sent_count

    async def _send_to_connection(self, conn_info: ConnectionInfo, message: dict[str, Any]) -> None:
        """Send a message to a specific connection.

        Args:
            conn_info: The connection info
            message: The message to send
        """
        await conn_info.websocket.send_json(message)
        conn_info.messages_sent += 1
        self._stats.total_messages_sent += 1

    async def _redis_subscription_loop(self, zone_id: str) -> None:
        """Subscribe to Redis events for a zone and broadcast to WebSocket clients.

        This creates a single Redis subscription per zone, shared by all
        WebSocket connections for that zone.

        Args:
            zone_id: The zone ID to subscribe to
        """
        if not self._event_bus:
            return

        logger.info(f"Starting Redis subscription for zone {zone_id}")

        try:
            async for event in self._event_bus.subscribe(zone_id):
                if not self._started:
                    break

                # Broadcast to all WebSocket clients for this zone
                sent = await self.broadcast_to_zone(zone_id, event)
                if sent > 0:
                    logger.debug(f"Broadcast {event.type} on {event.path} to {sent} clients")

        except asyncio.CancelledError:
            logger.debug(f"Redis subscription cancelled for zone {zone_id}")
            raise
        except Exception as e:
            logger.error(f"Redis subscription error for zone {zone_id}: {e}")

    async def _heartbeat_loop(self) -> None:
        """Send periodic pings to all connections and check for stale connections."""
        while self._started:
            try:
                await asyncio.sleep(PING_INTERVAL)

                if not self._started:
                    break

                now = time.time()
                stale_connections: list[str] = []

                # Check all connections
                for _zone_id, connections in list(self._connections.items()):
                    for conn_id, conn_info in list(connections.items()):
                        # Check if connection is stale (no pong received)
                        time_since_pong = now - conn_info.last_pong
                        if time_since_pong > PING_INTERVAL + PING_TIMEOUT:
                            logger.warning(
                                f"Connection {conn_id} stale (no pong in {time_since_pong:.1f}s)"
                            )
                            stale_connections.append(conn_id)
                            continue

                        # Send ping
                        try:
                            await self._send_to_connection(conn_info, {"type": "ping"})
                        except Exception as e:
                            logger.warning(f"Failed to ping {conn_id}: {e}")
                            stale_connections.append(conn_id)

                # Disconnect stale connections
                for conn_id in stale_connections:
                    with suppress(Exception):  # Best-effort close of stale connection
                        stale_zone_id = self._connection_to_zone.get(conn_id)
                        if stale_zone_id:
                            stale_conn = self._connections.get(stale_zone_id, {}).get(conn_id)
                            if stale_conn:
                                await stale_conn.websocket.close(
                                    code=1001, reason="Connection timeout"
                                )
                    await self.disconnect(conn_id)

            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Heartbeat loop error: {e}")
                await asyncio.sleep(5.0)  # Back off on errors

    def get_stats(self) -> dict[str, Any]:
        """Get WebSocket manager statistics.

        Returns:
            Statistics dictionary for health endpoint
        """
        return {
            "total_connections": self._stats.total_connections,
            "total_disconnections": self._stats.total_disconnections,
            "current_connections": self._stats.current_connections,
            "total_messages_sent": self._stats.total_messages_sent,
            "total_messages_received": self._stats.total_messages_received,
            "connections_by_zone": dict(self._stats.connections_by_zone),
            "active_zone_subscriptions": len(self._subscription_tasks),
        }

    def get_connection_count(self) -> int:
        """Get the current number of active connections.

        Returns:
            Number of active WebSocket connections
        """
        return self._stats.current_connections

    def get_connection_info(self, connection_id: str) -> ConnectionInfo | None:
        """Get information about a specific connection.

        Args:
            connection_id: The connection ID

        Returns:
            ConnectionInfo or None if not found
        """
        zone_id = self._connection_to_zone.get(connection_id)
        if not zone_id:
            return None
        return self._connections.get(zone_id, {}).get(connection_id)
