"""WebSocket Manager for real-time client notifications."""

# Standard library imports
import asyncio
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional, Set

# Third-party imports
from fastapi import WebSocket
import structlog

logger = structlog.get_logger(__name__)


class WebSocketManager:
    """WebSocket connection manager for sending notifications to clients.

    Supported events:
    - connection_expired: Connection expired
    - connection_deleted: Connection deleted manually
    - connection_will_expire: Connection will expire soon (5 min warning)
    - jwt_will_expire: JWT token will expire soon (5 min warning)
    - jwt_expired: JWT token expired
    - connection_extended: Connection TTL extended
    """

    def __init__(self) -> None:
        """Initialize WebSocket manager."""
        self.active_connections: Dict[str, Set[WebSocket]] = {}
        self._lock = asyncio.Lock()
    
    async def connect(self, websocket: WebSocket, username: str) -> None:
        """Connect a new client.

        Args:
            websocket: WebSocket connection (already accepted)
            username: Username
        """
        async with self._lock:
            if username not in self.active_connections:
                self.active_connections[username] = set()
            self.active_connections[username].add(websocket)

        logger.info(
            "WebSocket client connected",
            username=username,
            total_connections=len(
                self.active_connections.get(username, set())
            ),
        )

    async def disconnect(self, websocket: WebSocket, username: str) -> None:
        """Disconnect a client.

        Args:
            websocket: WebSocket connection
            username: Username
        """
        async with self._lock:
            if username in self.active_connections:
                self.active_connections[username].discard(websocket)

                if not self.active_connections[username]:
                    del self.active_connections[username]

        logger.info(
            "WebSocket client disconnected",
            username=username,
            remaining_connections=len(
                self.active_connections.get(username, set())
            ),
        )
    
    async def send_to_user(
        self, username: str, message: Dict[str, Any]
    ) -> None:
        """Send message to all WebSocket connections of a user.

        Args:
            username: Username
            message: Dictionary with data to send
        """
        if username not in self.active_connections:
            logger.debug(
                "No active WebSocket connections for user", username=username
            )
            return

        connections = self.active_connections[username].copy()

        disconnected = []
        for websocket in connections:
            try:
                await websocket.send_json(message)
                logger.debug(
                    "Message sent via WebSocket",
                    username=username,
                    event_type=message.get("type"),
                )
            except Exception as e:
                logger.error(
                    "Failed to send WebSocket message",
                    username=username,
                    error=str(e),
                )
                disconnected.append(websocket)

        if disconnected:
            async with self._lock:
                for ws in disconnected:
                    self.active_connections[username].discard(ws)

                if not self.active_connections[username]:
                    del self.active_connections[username]
    
    async def send_connection_expired(
        self,
        username: str,
        connection_id: str,
        hostname: str,
        protocol: str,
    ) -> None:
        """Notify about connection expiration.

        Args:
            username: Username
            connection_id: Connection ID
            hostname: Machine hostname
            protocol: Connection protocol
        """
        message = {
            "type": "connection_expired",
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "data": {
                "connection_id": connection_id,
                "hostname": hostname,
                "protocol": protocol,
                "reason": "TTL expired",
            },
        }
        await self.send_to_user(username, message)

        logger.info(
            "Connection expired notification sent",
            username=username,
            connection_id=connection_id,
            hostname=hostname,
        )
    
    async def send_connection_deleted(
        self,
        username: str,
        connection_id: str,
        hostname: str,
        protocol: str,
        reason: str = "manual",
    ) -> None:
        """Notify about connection deletion.

        Args:
            username: Username
            connection_id: Connection ID
            hostname: Machine hostname
            protocol: Connection protocol
            reason: Deletion reason (manual, expired, error)
        """
        message = {
            "type": "connection_deleted",
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "data": {
                "connection_id": connection_id,
                "hostname": hostname,
                "protocol": protocol,
                "reason": reason,
            },
        }
        await self.send_to_user(username, message)

        logger.info(
            "Connection deleted notification sent",
            username=username,
            connection_id=connection_id,
            reason=reason,
        )
    
    async def send_connection_will_expire(
        self,
        username: str,
        connection_id: str,
        hostname: str,
        protocol: str,
        minutes_remaining: int,
    ) -> None:
        """Warn about upcoming connection expiration.

        Args:
            username: Username
            connection_id: Connection ID
            hostname: Machine hostname
            protocol: Connection protocol
            minutes_remaining: Minutes until expiration
        """
        message = {
            "type": "connection_will_expire",
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "data": {
                "connection_id": connection_id,
                "hostname": hostname,
                "protocol": protocol,
                "minutes_remaining": minutes_remaining,
            },
        }
        await self.send_to_user(username, message)

        logger.info(
            "Connection expiration warning sent",
            username=username,
            connection_id=connection_id,
            minutes_remaining=minutes_remaining,
        )
    
    async def send_jwt_will_expire(
        self, username: str, minutes_remaining: int
    ) -> None:
        """Warn about upcoming JWT token expiration.

        Args:
            username: Username
            minutes_remaining: Minutes until expiration
        """
        message = {
            "type": "jwt_will_expire",
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "data": {
                "minutes_remaining": minutes_remaining,
                "action_required": "Please refresh your token or re-login",
            },
        }
        await self.send_to_user(username, message)

        logger.info(
            "JWT expiration warning sent",
            username=username,
            minutes_remaining=minutes_remaining,
        )

    async def send_jwt_expired(self, username: str) -> None:
        """Notify about JWT token expiration.

        Args:
            username: Username
        """
        message = {
            "type": "jwt_expired",
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "data": {"action_required": "Please re-login"},
        }
        await self.send_to_user(username, message)

        logger.info("JWT expired notification sent", username=username)
    
    async def send_connection_extended(
        self,
        username: str,
        connection_id: str,
        hostname: str,
        new_expires_at: datetime,
        additional_minutes: int,
    ) -> None:
        """Notify about connection extension.

        Args:
            username: Username
            connection_id: Connection ID
            hostname: Machine hostname
            new_expires_at: New expiration time
            additional_minutes: Minutes added
        """
        message = {
            "type": "connection_extended",
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "data": {
                "connection_id": connection_id,
                "hostname": hostname,
                "new_expires_at": new_expires_at.isoformat(),
                "additional_minutes": additional_minutes,
            },
        }
        await self.send_to_user(username, message)

        logger.info(
            "Connection extension notification sent",
            username=username,
            connection_id=connection_id,
            additional_minutes=additional_minutes,
        )

    def get_active_users(self) -> List[str]:
        """Get list of users with active WebSocket connections.

        Returns:
            List of usernames
        """
        return list(self.active_connections.keys())

    def get_connection_count(self, username: Optional[str] = None) -> int:
        """Get count of active WebSocket connections.

        Args:
            username: Username (if None, returns total count)

        Returns:
            Number of connections
        """
        if username:
            return len(self.active_connections.get(username, []))

        return sum(
            len(connections)
            for connections in self.active_connections.values()
        )


# Singleton instance
websocket_manager = WebSocketManager()

