"""WebSocket connection manager for the kommunikation plugin."""

from __future__ import annotations

import json
import logging
from typing import Any

from fastapi import WebSocket

logger = logging.getLogger(__name__)


class WebSocketManager:
    """Manages WebSocket connections per user for real-time messaging."""

    def __init__(self) -> None:
        # user_id (str) → list of WebSocket connections
        self._connections: dict[str, list[WebSocket]] = {}
        # conversation_id (str) → set of user_ids subscribed
        self._subscriptions: dict[str, set[str]] = {}

    async def connect(self, websocket: WebSocket, user_id: str) -> None:
        """Accept and register a new WebSocket connection."""
        await websocket.accept()
        if user_id not in self._connections:
            self._connections[user_id] = []
        self._connections[user_id].append(websocket)
        logger.debug(f"WebSocket connected: user={user_id}, total={len(self._connections[user_id])}")

    async def disconnect(self, websocket: WebSocket, user_id: str) -> None:
        """Remove a WebSocket connection."""
        conns = self._connections.get(user_id, [])
        if websocket in conns:
            conns.remove(websocket)
        if not conns:
            self._connections.pop(user_id, None)
            # Remove from all subscriptions
            for conv_id, users in self._subscriptions.items():
                users.discard(user_id)
        logger.debug(f"WebSocket disconnected: user={user_id}, remaining={len(conns)}")

    def subscribe(self, conversation_id: str, user_id: str) -> None:
        """Subscribe a user to a conversation's updates."""
        if conversation_id not in self._subscriptions:
            self._subscriptions[conversation_id] = set()
        self._subscriptions[conversation_id].add(user_id)

    def unsubscribe(self, conversation_id: str, user_id: str) -> None:
        """Unsubscribe a user from a conversation."""
        if conversation_id in self._subscriptions:
            self._subscriptions[conversation_id].discard(user_id)

    async def send_to_user(self, user_id: str, message: dict[str, Any]) -> None:
        """Send a message to all connections of a specific user."""
        conns = self._connections.get(user_id, [])
        text = json.dumps(message, default=str)
        for ws in conns:
            try:
                await ws.send_text(text)
            except Exception:
                logger.warning(f"Failed to send to user {user_id}, removing connection")
                await self.disconnect(ws, user_id)

    async def send_to_conversation(
        self,
        conversation_id: str,
        message: dict[str, Any],
        exclude_user: str | None = None,
    ) -> None:
        """Send a message to all users subscribed to a conversation."""
        user_ids = self._subscriptions.get(conversation_id, set())
        for user_id in list(user_ids):
            if exclude_user and user_id == exclude_user:
                continue
            await self.send_to_user(user_id, message)

    async def broadcast(self, message: dict[str, Any]) -> None:
        """Broadcast a message to all connected users."""
        for user_id in list(self._connections.keys()):
            await self.send_to_user(user_id, message)

    def get_online_users(self) -> list[str]:
        """Get list of currently connected user IDs."""
        return list(self._connections.keys())

    def is_user_online(self, user_id: str) -> bool:
        """Check if a user has any active connections."""
        return user_id in self._connections and len(self._connections[user_id]) > 0